fix(base-driver): track and remove BiDi subscriptions by ID (#22895)

packages/appium/docs/en/reference/api/bidi.md

@@ -46,7 +46,7 @@ Subscribes to one or more BiDi events.
 
 #### Response
 
-`null`
+`{subscription: string}` — a unique subscription ID that can be passed to `session.unsubscribe`.
 
 ### bidiUnsubscribe
 
@@ -56,15 +56,20 @@ session.unsubscribe
 
 > WebDriver BiDi documentation: [session.unsubscribe](https://w3c.github.io/webdriver-bidi/#command-session-unsubscribe)
 
-Unsubscribes from one or more BiDi events.
+Unsubscribes using either previously returned subscription IDs or event names.
+Removing a subscription ID preserves any overlapping subscriptions.
 
 #### Parameters
 
 |Name|Description|Type|Default|
 |--|--|--|--|
-|`contexts?`|Contexts in which to unsubscribe from the specified events. By default, the global scope is applied.|string[]|`['']`|
-|`events`|Names of events to unsubscribe from|string[]||
+|`subscriptions?`|Non-empty list of subscription IDs returned by `session.subscribe`. Cannot be combined with `events` or `contexts`.|string[]||
+|`contexts?`|Non-standard, only together with `events`: contexts in which to unsubscribe from the specified events. By default, the global scope is applied.|string[]|`['']`|
+|`events?`|Non-empty list of event names to unsubscribe from. Required when `subscriptions` is omitted.|string[]||
+
+All subscription IDs must be known to the current session. An unknown ID causes an `invalid argument`
+error without removing any subscriptions. Subscription IDs are cleared when the session is deleted.
 
 #### Response
 
-`null`
+`{}`

packages/base-driver/lib/basedriver/commands/bidi-subscriptions.ts

@@ -0,0 +1,36 @@
+import {util} from '@appium/support';
+
+type BidiSubscriptionDriver = {bidiEventSubs: Record<string, string[]>};
+
+// subscription id -> event name -> contexts; kept per id so removing one keeps overlapping ones
+type Subscriptions = Map<string, Map<string, string[]>>;
+
+const subscriptionsByDriver = new WeakMap<object, Subscriptions>();
+
+/** Get (or lazily create) the subscription registry of the given driver */
+export function getSubscriptions(driver: object): Subscriptions {
+  let subscriptions = subscriptionsByDriver.get(driver);
+  if (!subscriptions) {
+    subscriptions = new Map();
+    subscriptionsByDriver.set(driver, subscriptions);
+  }
+  return subscriptions;
+}
+
+/** Recompute `bidiEventSubs` as the union of all registered subscriptions */
+export function refreshEventSubscriptions(driver: BidiSubscriptionDriver): void {
+  const events = new Map<string, string[]>();
+  for (const subscription of getSubscriptions(driver).values()) {
+    for (const [event, contexts] of subscription) {
+      events.set(event, util.uniq([...(events.get(event) ?? []), ...contexts]));
+    }
+  }
+  // Object.fromEntries defines own properties, so names like '__proto__' stay safe
+  driver.bidiEventSubs = Object.fromEntries(events);
+}
+
+/** Forget all subscriptions, e.g. when the session ends */
+export function clearBidiSubscriptions(driver: BidiSubscriptionDriver): void {
+  subscriptionsByDriver.delete(driver);
+  driver.bidiEventSubs = {};
+}

packages/base-driver/lib/basedriver/commands/bidi.ts

@@ -1,7 +1,9 @@
 import {util} from '@appium/support';
 import type {Constraints, DriverStatus, IBidiCommands} from '@appium/types';
 
+import {errors} from '../../protocol/errors';
 import type {BaseDriver} from '../driver';
+import {getSubscriptions, refreshEventSubscriptions} from './bidi-subscriptions';
 import {mixin} from './mixin';
 
 declare module '../driver' {
@@ -9,25 +11,67 @@ declare module '../driver' {
   interface BaseDriver<C extends Constraints> extends IBidiCommands {}
 }
 
+function assertStringList(value: unknown, name: string): asserts value is string[] {
+  if (!Array.isArray(value) || !value.length || value.some((item) => typeof item !== 'string' || !item.length)) {
+    throw new errors.InvalidArgumentError(`${name} must be a non-empty array of non-empty strings`);
+  }
+}
+
+function assertContexts(value: unknown): asserts value is string[] {
+  if (!Array.isArray(value) || value.some((context) => typeof context !== 'string')) {
+    throw new errors.InvalidArgumentError('contexts must be an array of strings');
+  }
+}
+
 const BidiCommands: IBidiCommands = {
   async bidiSubscribe<C extends Constraints>(this: BaseDriver<C>, events: string[], contexts: string[] = ['']) {
-    for (const event of events) {
-      // a later subscribe for the same event must keep existing contexts; the spec appends subscriptions
-      const existing = this.bidiEventSubs[event] ?? [];
-      this.bidiEventSubs[event] = util.uniq([...existing, ...contexts]);
-    }
+    assertStringList(events, 'events');
+    assertContexts(contexts);
+    const subscription = util.uuidV4();
+    getSubscriptions(this).set(subscription, new Map(events.map((event) => [event, [...contexts]])));
+    refreshEventSubscriptions(this);
+    return {subscription};
   },
 
-  async bidiUnsubscribe<C extends Constraints>(this: BaseDriver<C>, events: string[], contexts: string[] = ['']) {
-    for (const event of events) {
-      if (!this.bidiEventSubs[event]) {
-        continue;
+  async bidiUnsubscribe<C extends Constraints>(
+    this: BaseDriver<C>,
+    events?: string[],
+    contexts?: string[],
+    subscriptions?: string[],
+  ) {
+    const registered = getSubscriptions(this);
+    if (subscriptions === undefined) {
+      const contextsToRemove = contexts ?? [''];
+      assertStringList(events, 'events');
+      assertContexts(contextsToRemove);
+      for (const subscription of registered.values()) {
+        for (const event of events) {
+          const remaining = subscription.get(event)?.filter((c) => !contextsToRemove.includes(c));
+          if (remaining?.length) {
+            subscription.set(event, remaining);
+          } else {
+            subscription.delete(event);
+          }
+        }
       }
-      this.bidiEventSubs[event] = this.bidiEventSubs[event].filter((c) => !contexts.includes(c));
-      if (this.bidiEventSubs[event].length === 0) {
-        delete this.bidiEventSubs[event];
+      refreshEventSubscriptions(this);
+      return;
+    }
+
+    assertStringList(subscriptions, 'subscriptions');
+    if (events !== undefined || contexts !== undefined) {
+      throw new errors.InvalidArgumentError('subscriptions cannot be combined with events or contexts');
+    }
+    // validate everything before removing anything
+    for (const id of subscriptions) {
+      if (!registered.has(id)) {
+        throw new errors.InvalidArgumentError(`Unknown subscription id: ${id}`);
       }
     }
+    for (const id of subscriptions) {
+      registered.delete(id);
+    }
+    refreshEventSubscriptions(this);
   },
 
   async bidiStatus<C extends Constraints>(this: BaseDriver<C>): Promise<DriverStatus> {

packages/base-driver/lib/basedriver/driver.ts

@@ -26,6 +26,7 @@ import {calcSignature} from '../helpers/session';
 import {DELETE_SESSION_COMMAND, determineProtocol, errors} from '../protocol';
 import {mergePlainObjects} from '../utils';
 import {processCapabilities, validateCaps} from './capabilities';
+import {clearBidiSubscriptions} from './commands/bidi-subscriptions';
 import {DriverCore} from './core';
 import * as helpers from './helpers';
 
@@ -392,6 +393,7 @@ export class BaseDriver<
       }
     }
     this.sessionId = null;
+    clearBidiSubscriptions(this);
   }
 
   logExtraCaps(caps: Capabilities<C>) {

packages/base-driver/lib/protocol/bidi-commands.ts

@@ -1,19 +1,25 @@
 import type {BidiModuleMap} from '@appium/types';
 
-const SUBSCRIPTION_REQUEST_PARAMS = {
+const SUBSCRIBE_PARAMS = {
   required: ['events'],
   optional: ['contexts'],
 } as const;
 
+// `subscriptions` and `events` are alternative spec forms; `contexts` is a non-standard
+// addition to `events`. `subscriptions` excludes the others, which bidiUnsubscribe validates.
+const UNSUBSCRIBE_PARAMS = {
+  optional: ['events', 'contexts', 'subscriptions'],
+} as const;
+
 export const BIDI_COMMANDS = {
   session: {
     subscribe: {
       command: 'bidiSubscribe',
-      params: SUBSCRIPTION_REQUEST_PARAMS,
+      params: SUBSCRIBE_PARAMS,
     },
     unsubscribe: {
       command: 'bidiUnsubscribe',
-      params: SUBSCRIPTION_REQUEST_PARAMS,
+      params: UNSUBSCRIBE_PARAMS,
     },
     status: {
       command: 'bidiStatus',

packages/base-driver/test/unit/basedriver/commands/bidi.spec.ts

@@ -13,6 +13,98 @@ describe('bidi commands -', function () {
   });
 
   describe('bidiUnsubscribe', function () {
+    it('should return unique IDs and remove only the selected overlapping subscription', async function () {
+      const first = await driver.executeBidiCommand('session.subscribe', {events: ['log.entryAdded'], contexts: ['a']});
+      const second = await driver.executeBidiCommand('session.subscribe', {
+        events: ['log.entryAdded'],
+        contexts: ['a', 'b'],
+      });
+      assert.equal(typeof first.subscription, 'string');
+      assert.notEqual(first.subscription, second.subscription);
+      assert.deepEqual(
+        await driver.executeBidiCommand('session.unsubscribe', {subscriptions: [first.subscription]}),
+        {},
+      );
+      assert.deepEqual(driver.bidiEventSubs, {'log.entryAdded': ['a', 'b']});
+      await driver.executeBidiCommand('session.unsubscribe', {subscriptions: [second.subscription]});
+      assert.deepEqual(driver.bidiEventSubs, {});
+    });
+
+    it('should validate every ID before removing any subscriptions', async function () {
+      const {subscription} = await driver.bidiSubscribe(['log.entryAdded']);
+      await assert.rejects(
+        driver.executeBidiCommand('session.unsubscribe', {subscriptions: [subscription, 'missing']}),
+        /Unknown subscription/,
+      );
+      assert.deepEqual(driver.bidiEventSubs, {'log.entryAdded': ['']});
+      await driver.executeBidiCommand('session.unsubscribe', {subscriptions: [subscription, subscription]});
+      assert.deepEqual(driver.bidiEventSubs, {});
+    });
+
+    it('should support unsubscribing by event names without confusing events and IDs', async function () {
+      const {subscription} = await driver.bidiSubscribe(['log.entryAdded'], ['a', 'b']);
+      await driver.executeBidiCommand('session.unsubscribe', {events: ['log.entryAdded'], contexts: ['a']});
+      assert.deepEqual(driver.bidiEventSubs, {'log.entryAdded': ['b']});
+      await driver.executeBidiCommand('session.unsubscribe', {subscriptions: [subscription]});
+      assert.deepEqual(driver.bidiEventSubs, {});
+    });
+
+    it('should reject mixed ID and event requests before mutating subscriptions', async function () {
+      const {subscription} = await driver.bidiSubscribe(['log.entryAdded']);
+      for (const extra of [{events: ['log.entryAdded']}, {contexts: ['a']}, {contexts: []}]) {
+        await assert.rejects(
+          driver.executeBidiCommand('session.unsubscribe', {subscriptions: [subscription], ...extra}),
+          /cannot be combined/,
+        );
+      }
+      assert.deepEqual(driver.bidiEventSubs, {'log.entryAdded': ['']});
+    });
+
+    it('should reject mixed arguments in direct ID-based unsubscribe calls without removing subscriptions', async function () {
+      const {subscription} = await driver.bidiSubscribe(['log.entryAdded']);
+      for (const contexts of [[], [''], ['a']]) {
+        await assert.rejects(driver.bidiUnsubscribe(undefined, contexts, [subscription]), /cannot be combined/);
+      }
+      await assert.rejects(driver.bidiUnsubscribe(['log.entryAdded'], undefined, [subscription]), /cannot be combined/);
+      assert.deepEqual(driver.bidiEventSubs, {'log.entryAdded': ['']});
+      await driver.bidiUnsubscribe(undefined, undefined, [subscription]);
+      assert.deepEqual(driver.bidiEventSubs, {});
+    });
+
+    it('should default event-based unsubscribe to the global context', async function () {
+      await driver.bidiSubscribe(['log.entryAdded']);
+      await driver.bidiUnsubscribe(['log.entryAdded']);
+      assert.deepEqual(driver.bidiEventSubs, {});
+    });
+
+    it('should safely aggregate event names matching Object prototype properties', async function () {
+      const events = ['__proto__', 'constructor', 'hasOwnProperty'];
+      const first = await driver.bidiSubscribe(events, ['a']);
+      const second = await driver.bidiSubscribe(events, ['a', 'b']);
+      assert.deepEqual(driver.bidiEventSubs, Object.fromEntries(events.map((event) => [event, ['a', 'b']])));
+      assert.equal(Object.getPrototypeOf(driver.bidiEventSubs), Object.prototype);
+      await driver.bidiUnsubscribe(undefined, undefined, [second.subscription]);
+      assert.deepEqual(driver.bidiEventSubs, Object.fromEntries(events.map((event) => [event, ['a']])));
+      await driver.bidiUnsubscribe(undefined, undefined, [first.subscription]);
+      assert.deepEqual(driver.bidiEventSubs, {});
+    });
+
+    it('should clear subscription IDs when the session is deleted', async function () {
+      const {subscription} = await driver.bidiSubscribe(['log.entryAdded']);
+      await driver.deleteSession();
+      assert.deepEqual(driver.bidiEventSubs, {});
+      await assert.rejects(
+        driver.executeBidiCommand('session.unsubscribe', {subscriptions: [subscription]}),
+        /Unknown subscription/,
+      );
+    });
+
+    for (const subscriptions of [[], [''], 'id', [null]]) {
+      it(`should reject malformed subscription IDs: ${JSON.stringify(subscriptions)}`, async function () {
+        await assert.rejects(driver.executeBidiCommand('session.unsubscribe', {subscriptions}), /non-empty array/);
+      });
+    }
+
     it('should not throw when the event was never subscribed', async function () {
       await driver.bidiUnsubscribe(['log.entryAdded'], ['']);
       assert.deepStrictEqual(driver.bidiEventSubs, {});

packages/types/lib/commands/basedriver.ts

@@ -3,8 +3,8 @@ import type {Constraints} from '../constraints';
 import type {Element, StringRecord} from '../util';
 
 export interface IBidiCommands {
-  bidiSubscribe(events: string[], contexts: string[]): Promise<void>;
-  bidiUnsubscribe(events: string[], contexts: string[]): Promise<void>;
+  bidiSubscribe(events: string[], contexts?: string[]): Promise<{subscription: string}>;
+  bidiUnsubscribe(events?: string[], contexts?: string[], subscriptions?: string[]): Promise<void>;
   bidiStatus(): Promise<DriverStatus>;
 }