stash

packages/composable-controller/CHANGELOG.md

@@ -16,6 +16,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - **BREAKING:** Passing a non-controller into `controllers` constructor option now throws an error ([#3904](https://github.com/MetaMask/core/pull/3904))
 - **BREAKING:** The `AllowedActions` parameter of the `ComposableControllerMessenger` type is narrowed from `string` to `never`, as `ComposableController` does not use any external controller actions. ([#3904](https://github.com/MetaMask/core/pull/3904))
+- **BREAKING:** The `ComposableController` class is now a generic class that expects one generic argument `ChildControllers` in the form of a union of the controllers passed into the `controllers` array constructor option ([#3941](https://github.com/MetaMask/core/pull/3941))
+  - Child controllers that extend `BaseControllerV1` must have an overridden `name` property that is defined using the `as const` assertion for the `ComposableController` class to be typed correctly.
+  - The controller state is now constrained by a `ComposedControllerState` generic argument, which is automaticallly derived from the `ChildControllers` argument without needing to be passed in by the user.
+  - The `messenger` constructor option is constrained by a `ComposedControllerMessenger` generic argument which is also automatically derived, and is required to contain the `stateChange` events for all child controllers in its `Events` and `AllowedEvents` parameters, as well as the composable controller instance in its `Events` parameter.
 - Add `@metamask/utils` ^8.3.0 as a dependency. ([#3904](https://github.com/MetaMask/core/pull/3904))
 
 ### Removed

packages/composable-controller/src/ComposableController.test.ts

@@ -60,6 +60,50 @@ class FooController extends BaseController<
   }
 }
 
+type QuzControllerState = {
+  quz: string;
+};
+type QuzControllerEvent = {
+  type: `QuzController:stateChange`;
+  payload: [QuzControllerState, Patch[]];
+};
+
+type QuzMessenger = RestrictedControllerMessenger<
+  'QuzController',
+  never,
+  QuzControllerEvent,
+  never,
+  never
+>;
+
+const quzControllerStateMetadata = {
+  quz: {
+    persist: true,
+    anonymous: true,
+  },
+};
+
+class QuzController extends BaseController<
+  'QuzController',
+  QuzControllerState,
+  QuzMessenger
+> {
+  constructor(messagingSystem: QuzMessenger) {
+    super({
+      messenger: messagingSystem,
+      metadata: quzControllerStateMetadata,
+      name: 'QuzController',
+      state: { quz: 'quz' },
+    });
+  }
+
+  updateQuz(quz: string) {
+    super.update((state) => {
+      state.quz = quz;
+    });
+  }
+}
+
 // Mock BaseControllerV1 classes
 
 type BarControllerState = BaseState & {
@@ -71,7 +115,7 @@ class BarController extends BaseControllerV1<never, BarControllerState> {
     bar: 'bar',
   };
 
-  override name = 'BarController';
+  override name = 'BarController' as const;
 
   constructor() {
     super();
@@ -92,7 +136,7 @@ class BazController extends BaseControllerV1<never, BazControllerState> {
     baz: 'baz',
   };
 
-  override name = 'BazController';
+  override name = 'BazController' as const;
 
   constructor() {
     super();
@@ -157,23 +201,39 @@ describe('ComposableController', () => {
     it('should compose controller state', () => {
       const controllerMessenger = new ControllerMessenger<
         never,
-        FooControllerEvent
+        FooControllerEvent | QuzControllerEvent
       >();
-      const fooControllerMessenger = controllerMessenger.getRestricted({
+      const fooMessenger = controllerMessenger.getRestricted<
+        'FooController',
+        never,
+        never
+      >({
         name: 'FooController',
       });
-      const fooController = new FooController(fooControllerMessenger);
+      const quzMessenger = controllerMessenger.getRestricted<
+        'QuzController',
+        never,
+        never
+      >({
+        name: 'QuzController',
+      });
+      const fooController = new FooController(fooMessenger);
+      const quzController = new QuzController(quzMessenger);
 
       const composableControllerMessenger = controllerMessenger.getRestricted({
         name: 'ComposableController',
-        allowedEvents: ['FooController:stateChange'],
+        allowedEvents: [
+          'FooController:stateChange',
+          'QuzController:stateChange',
+        ],
       });
       const composableController = new ComposableController({
-        controllers: [fooController],
+        controllers: [fooController, quzController],
         messenger: composableControllerMessenger,
       });
       expect(composableController.state).toStrictEqual({
         FooController: { foo: 'foo' },
+        QuzController: { quz: 'quz' },
       });
     });
 

packages/composable-controller/src/ComposableController.ts

@@ -4,6 +4,7 @@ import type {
   RestrictedControllerMessenger,
   BaseState,
   BaseConfig,
+  EventConstraint,
   StateMetadata,
 } from '@metamask/base-controller';
 import { isValidJson, type Json } from '@metamask/utils';
@@ -112,13 +113,43 @@ export type ComposableControllerMessenger = RestrictedControllerMessenger<
   AllowedEvents['type']
 >;
 
+type GetStateChangeEventsUnion<Controller extends ControllerInstance> =
+  Controller extends ControllerInstance
+    ? Controller extends BaseControllerV2Instance
+      ? ControllerStateChangeEvent<Controller['name'], Controller['state']>
+      : never
+    : never;
+
 /**
  * Controller that can be used to compose multiple controllers together.
  */
-export class ComposableController extends BaseController<
+export class ComposableController<
+  ChildControllers extends ControllerInstance,
+  ComposedControllerState extends ComposableControllerState = {
+    [P in ChildControllers as P['name']]: P extends ControllerInstance
+      ? P['state']
+      : never;
+  },
+  ComposedControllerStateChangeEvent extends EventConstraint & {
+    type: `${typeof controllerName}:stateChange`;
+  } = ControllerStateChangeEvent<
+    typeof controllerName,
+    ComposedControllerState
+  >,
+  ChildControllersStateChangeEvents extends EventConstraint & {
+    type: `${string}:stateChange`;
+  } = GetStateChangeEventsUnion<ChildControllers>,
+  ComposedControllerMessenger extends ComposableControllerMessenger = RestrictedControllerMessenger<
+    typeof controllerName,
+    never,
+    ComposedControllerStateChangeEvent | ChildControllersStateChangeEvents,
+    never,
+    ChildControllersStateChangeEvents['type']
+  >,
+> extends BaseController<
   typeof controllerName,
-  ComposableControllerState,
-  ComposableControllerMessenger
+  ComposedControllerState,
+  ComposedControllerMessenger
 > {
   /**
    * Creates a ComposableController instance.
@@ -132,29 +163,29 @@ export class ComposableController extends BaseController<
     controllers,
     messenger,
   }: {
-    controllers: ControllerInstance[];
-    messenger: ComposableControllerMessenger;
+    controllers: ChildControllers[];
+    messenger: ComposedControllerMessenger;
   }) {
     if (messenger === undefined) {
       throw new Error(`Messaging system is required`);
     }
 
     super({
       name: controllerName,
-      metadata: controllers.reduce<StateMetadata<ComposableControllerState>>(
+      metadata: controllers.reduce<StateMetadata<ComposedControllerState>>(
         (metadata, controller) => ({
           ...metadata,
           [controller.name]: isBaseController(controller)
             ? controller.metadata
             : { persist: true, anonymous: true },
         }),
-        {},
+        {} as never,
       ),
-      state: controllers.reduce<ComposableControllerState>(
+      state: controllers.reduce<ComposedControllerState>(
         (state, controller) => {
           return { ...state, [controller.name]: controller.state };
         },
-        {},
+        {} as never,
       ),
       messenger,
     });
@@ -173,18 +204,24 @@ export class ComposableController extends BaseController<
     const { name } = controller;
     if (isBaseControllerV1(controller)) {
       controller.subscribe((childState) => {
-        this.update((state) => ({
-          ...state,
-          [name]: childState,
-        }));
+        this.update(
+          (state) =>
+            ({
+              ...state,
+              [name]: childState,
+            } as ComposedControllerState),
+        );
       });
     } else if (isBaseController(controller)) {
       this.messagingSystem.subscribe(`${name}:stateChange`, (childState) => {
         if (isValidJson(childState)) {
-          this.update((state) => ({
-            ...state,
-            [name]: childState,
-          }));
+          this.update(
+            (state) =>
+              ({
+                ...state,
+                [name]: childState,
+              } as ComposedControllerState),
+          );
         }
       });
     } else {