feat: implement user context subscriptions (GoogleChromeLabs#3039)

- Introduces UserContextStorage to abstract and re-use operations on user contexts. For now, only the get operation is moved. Issue GoogleChromeLabs#2983
anshikavashistha · Jan 31, 2025 · 33c49cb · 33c49cb
1 parent cb16a32
commit 33c49cb
Show file tree

Hide file tree

Showing 10 changed files with 342 additions and 62 deletions.
diff --git a/src/bidiMapper/BidiServer.ts b/src/bidiMapper/BidiServer.ts
@@ -27,6 +27,7 @@ import type {BidiCommandParameterParser} from './BidiParser.js';
 import type {BidiTransport} from './BidiTransport.js';
 import {CommandProcessor, CommandProcessorEvents} from './CommandProcessor.js';
 import {BluetoothProcessor} from './modules/bluetooth/BluetoothProcessor.js';
+import {UserContextStorage} from './modules/browser/UserContextStorage.js';
 import {CdpTargetManager} from './modules/cdp/CdpTargetManager.js';
 import {BrowsingContextStorage} from './modules/context/BrowsingContextStorage.js';
 import {NetworkStorage} from './modules/network/NetworkStorage.js';
@@ -91,7 +92,11 @@ export class BidiServer extends EventEmitter<BidiServerEvent> {
     );
     this.#transport = bidiTransport;
     this.#transport.setOnMessage(this.#handleIncomingMessage);
-    this.#eventManager = new EventManager(this.#browsingContextStorage);
+    const userUserContextStorage = new UserContextStorage(browserCdpClient);
+    this.#eventManager = new EventManager(
+      this.#browsingContextStorage,
+      userUserContextStorage,
+    );
     const networkStorage = new NetworkStorage(
       this.#eventManager,
       this.#browsingContextStorage,
@@ -111,6 +116,7 @@ export class BidiServer extends EventEmitter<BidiServerEvent> {
       this.#preloadScriptStorage,
       networkStorage,
       this.#bluetoothProcessor,
+      userUserContextStorage,
       parser,
       async (options: MapperOptions) => {
         // This is required to ignore certificate errors when service worker is fetched.

diff --git a/src/bidiMapper/CommandProcessor.ts b/src/bidiMapper/CommandProcessor.ts
@@ -33,6 +33,7 @@ import type {BidiCommandParameterParser} from './BidiParser.js';
 import type {MapperOptions} from './BidiServer.js';
 import type {BluetoothProcessor} from './modules/bluetooth/BluetoothProcessor.js';
 import {BrowserProcessor} from './modules/browser/BrowserProcessor.js';
+import type {UserContextStorage} from './modules/browser/UserContextStorage.js';
 import {CdpProcessor} from './modules/cdp/CdpProcessor.js';
 import {BrowsingContextProcessor} from './modules/context/BrowsingContextProcessor.js';
 import type {BrowsingContextStorage} from './modules/context/BrowsingContextStorage.js';
@@ -85,6 +86,7 @@ export class CommandProcessor extends EventEmitter<CommandProcessorEventsMap> {
     preloadScriptStorage: PreloadScriptStorage,
     networkStorage: NetworkStorage,
     bluetoothProcessor: BluetoothProcessor,
+    userContextStorage: UserContextStorage,
     parser: BidiCommandParameterParser = new BidiNoOpParser(),
     initConnection: (options: MapperOptions) => Promise<void>,
     logger?: LoggerFn,
@@ -99,6 +101,7 @@ export class CommandProcessor extends EventEmitter<CommandProcessorEventsMap> {
     this.#browserProcessor = new BrowserProcessor(
       browserCdpClient,
       browsingContextStorage,
+      userContextStorage,
     );
     this.#browsingContextProcessor = new BrowsingContextProcessor(
       browserCdpClient,

diff --git a/src/bidiMapper/modules/browser/BrowserProcessor.ts b/src/bidiMapper/modules/browser/BrowserProcessor.ts
@@ -26,16 +26,21 @@ import {
 import type {CdpClient} from '../../BidiMapper.js';
 import type {BrowsingContextStorage} from '../context/BrowsingContextStorage.js';
 
+import type {UserContextStorage} from './UserContextStorage.js';
+
 export class BrowserProcessor {
   readonly #browserCdpClient: CdpClient;
   readonly #browsingContextStorage: BrowsingContextStorage;
+  readonly #userContextStorage: UserContextStorage;
 
   constructor(
     browserCdpClient: CdpClient,
     browsingContextStorage: BrowsingContextStorage,
+    userContextStorage: UserContextStorage,
   ) {
     this.#browserCdpClient = browserCdpClient;
     this.#browsingContextStorage = browsingContextStorage;
+    this.#userContextStorage = userContextStorage;
   }
 
   close(): EmptyResult {
@@ -90,20 +95,8 @@ export class BrowserProcessor {
   }
 
   async getUserContexts(): Promise<Browser.GetUserContextsResult> {
-    const result = await this.#browserCdpClient.sendCommand(
-      'Target.getBrowserContexts',
-    );
     return {
-      userContexts: [
-        {
-          userContext: 'default',
-        },
-        ...result.browserContextIds.map((id) => {
-          return {
-            userContext: id,
-          };
-        }),
-      ],
+      userContexts: await this.#userContextStorage.getUserContexts(),
     };
   }
 

diff --git a/src/bidiMapper/modules/browser/UserContextStorage.ts b/src/bidiMapper/modules/browser/UserContextStorage.ts
@@ -0,0 +1,45 @@
+/**
+ * Copyright 2025 Google LLC.
+ * Copyright (c) Microsoft Corporation.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import type {CdpClient} from '../../../cdp/CdpClient.js';
+import type {Browser} from '../../../protocol/protocol.js';
+
+export class UserContextStorage {
+  #browserClient: CdpClient;
+
+  constructor(browserClient: CdpClient) {
+    this.#browserClient = browserClient;
+  }
+
+  async getUserContexts(): Promise<
+    [Browser.UserContextInfo, ...Browser.UserContextInfo[]]
+  > {
+    const result = await this.#browserClient.sendCommand(
+      'Target.getBrowserContexts',
+    );
+    return [
+      {
+        userContext: 'default',
+      },
+      ...result.browserContextIds.map((id) => {
+        return {
+          userContext: id,
+        };
+      }),
+    ];
+  }
+}
diff --git a/src/bidiMapper/modules/network/NetworkStorage.spec.ts b/src/bidiMapper/modules/network/NetworkStorage.spec.ts
@@ -21,6 +21,7 @@ import type {CdpClient} from '../../../cdp/CdpClient.js';
 import {ChromiumBidi, Network} from '../../../protocol/protocol.js';
 import {ProcessingQueue} from '../../../utils/ProcessingQueue.js';
 import type {OutgoingMessage} from '../../OutgoingMessage.js';
+import {UserContextStorage} from '../browser/UserContextStorage.js';
 import type {CdpTarget} from '../cdp/CdpTarget.js';
 import type {BrowsingContextImpl} from '../context/BrowsingContextImpl.js';
 import {BrowsingContextStorage} from '../context/BrowsingContextStorage.js';
@@ -75,9 +76,10 @@ describe('NetworkStorage', () => {
       id: MockCdpNetworkEvents.defaultFrameId,
     } as unknown as BrowsingContextImpl;
     cdpClient = cdpTarget.cdpClient;
+    const userContextStorage = new UserContextStorage(cdpClient);
     // We need to add it the storage to emit properly
     browsingContextStorage.addContext(browsingContext);
-    eventManager = new EventManager(browsingContextStorage);
+    eventManager = new EventManager(browsingContextStorage, userContextStorage);
     processingQueue = new ProcessingQueue<OutgoingMessage>(
       async ({message}) => {
         if (message.type === 'event') {
@@ -93,6 +95,7 @@ describe('NetworkStorage', () => {
       // Verify that the Request send the message
       // To the correct context
       [MockCdpNetworkEvents.defaultFrameId],
+      [],
       {},
     );
     eventManager.on(EventManagerEvents.Event, ({message, event}) => {

diff --git a/src/bidiMapper/modules/session/EventManager.ts b/src/bidiMapper/modules/session/EventManager.ts
@@ -17,8 +17,10 @@
 
 import type {BidiPlusChannel} from '../../../protocol/chromium-bidi.js';
 import {
+  type Browser,
   ChromiumBidi,
   InvalidArgumentException,
+  NoSuchUserContextException,
   type BrowsingContext,
 } from '../../../protocol/protocol.js';
 import {Buffer} from '../../../utils/Buffer.js';
@@ -27,6 +29,7 @@ import {EventEmitter} from '../../../utils/EventEmitter.js';
 import {IdWrapper} from '../../../utils/IdWrapper.js';
 import type {Result} from '../../../utils/result.js';
 import {OutgoingMessage} from '../../OutgoingMessage.js';
+import type {UserContextStorage} from '../browser/UserContextStorage.js';
 import type {BrowsingContextStorage} from '../context/BrowsingContextStorage.js';
 
 import {assertSupportedEvent} from './events.js';
@@ -118,9 +121,15 @@ export class EventManager extends EventEmitter<EventManagerEventsMap> {
     ((contextId: BrowsingContext.BrowsingContext) => void)[]
   >;
 
-  constructor(browsingContextStorage: BrowsingContextStorage) {
+  #userContextStorage: UserContextStorage;
+
+  constructor(
+    browsingContextStorage: BrowsingContextStorage,
+    userContextStorage: UserContextStorage,
+  ) {
     super();
     this.#browsingContextStorage = browsingContextStorage;
+    this.#userContextStorage = userContextStorage;
     this.#subscriptionManager = new SubscriptionManager(browsingContextStorage);
     this.#subscribeHooks = new DefaultMap(() => []);
   }
@@ -213,12 +222,19 @@ export class EventManager extends EventEmitter<EventManagerEventsMap> {
   async subscribe(
     eventNames: ChromiumBidi.EventNames[],
     contextIds: BrowsingContext.BrowsingContext[],
+    userContextIds: Browser.UserContext[],
     channel: BidiPlusChannel,
   ): Promise<string> {
     for (const name of eventNames) {
       assertSupportedEvent(name);
     }
 
+    if (userContextIds.length && contextIds.length) {
+      throw new InvalidArgumentException(
+        'Both userContexts and contexts cannot be specified.',
+      );
+    }
+
     // First check if all the contexts are known.
     for (const contextId of contextIds) {
       if (contextId !== null) {
@@ -227,6 +243,21 @@ export class EventManager extends EventEmitter<EventManagerEventsMap> {
       }
     }
 
+    // Validate user contexts.
+    if (userContextIds.length) {
+      const userContexts = await this.#userContextStorage.getUserContexts();
+      const knownUserContextIds = new Set(
+        userContexts.map((userContext) => userContext.userContext),
+      );
+      for (const userContextId of userContextIds) {
+        if (!knownUserContextIds.has(userContextId)) {
+          throw new NoSuchUserContextException(
+            `User context ${userContextId} not found`,
+          );
+        }
+      }
+    }
+
     const unrolledEventNames = new Set(unrollEvents(eventNames));
     const subscribeStepEvents = new Map<ChromiumBidi.EventNames, Set<string>>();
     const subscriptionNavigableIds = new Set(
@@ -260,6 +291,7 @@ export class EventManager extends EventEmitter<EventManagerEventsMap> {
     const subscription = this.#subscriptionManager.subscribe(
       eventNames,
       contextIds,
+      userContextIds,
       channel,
     );
 

diff --git a/src/bidiMapper/modules/session/SessionProcessor.ts b/src/bidiMapper/modules/session/SessionProcessor.ts
@@ -149,6 +149,7 @@ export class SessionProcessor {
     const subscription = await this.#eventManager.subscribe(
       params.events as ChromiumBidi.EventNames[],
       params.contexts ?? [],
+      params.userContexts ?? [],
       channel,
     );
     return {