feat: Allow updating interface context (#2809)

Allows updating the `context` object in the `SnapInterfaceController`
via `snap_updateInterface`. This is a non-breaking change as `context`
is optional.

Closes #2804

Author: FrederikBolding

packages/snaps-controllers/coverage.json

@@ -1,5 +1,5 @@
 {
-  "branches": 92.6,
+  "branches": 92.62,
   "functions": 96.65,
   "lines": 97.97,
   "statements": 97.67

packages/snaps-controllers/src/interface/SnapInterfaceController.test.tsx

@@ -545,6 +545,111 @@ describe('SnapInterfaceController', () => {
       expect(state).toStrictEqual({ foo: { baz: null } });
     });
 
+    it('can update an interface and context', async () => {
+      const rootMessenger = getRootSnapInterfaceControllerMessenger();
+      const controllerMessenger =
+        getRestrictedSnapInterfaceControllerMessenger(rootMessenger);
+
+      /* eslint-disable-next-line no-new */
+      new SnapInterfaceController({
+        messenger: controllerMessenger,
+      });
+
+      const components = form({
+        name: 'foo',
+        children: [input({ name: 'bar' })],
+      });
+
+      const newContent = form({
+        name: 'foo',
+        children: [input({ name: 'baz' })],
+      });
+
+      const context = { foo: 'bar' };
+
+      const id = await rootMessenger.call(
+        'SnapInterfaceController:createInterface',
+        MOCK_SNAP_ID,
+        components,
+        context,
+      );
+
+      const newContext = { foo: 'baz' };
+
+      await rootMessenger.call(
+        'SnapInterfaceController:updateInterface',
+        MOCK_SNAP_ID,
+        id,
+        newContent,
+        newContext,
+      );
+
+      const {
+        content,
+        state,
+        context: interfaceContext,
+      } = rootMessenger.call(
+        'SnapInterfaceController:getInterface',
+        MOCK_SNAP_ID,
+        id,
+      );
+
+      expect(content).toStrictEqual(getJsxElementFromComponent(newContent));
+      expect(state).toStrictEqual({ foo: { baz: null } });
+      expect(interfaceContext).toStrictEqual(newContext);
+    });
+
+    it('does not replace context if none is provided', async () => {
+      const rootMessenger = getRootSnapInterfaceControllerMessenger();
+      const controllerMessenger =
+        getRestrictedSnapInterfaceControllerMessenger(rootMessenger);
+
+      /* eslint-disable-next-line no-new */
+      new SnapInterfaceController({
+        messenger: controllerMessenger,
+      });
+
+      const components = form({
+        name: 'foo',
+        children: [input({ name: 'bar' })],
+      });
+
+      const newContent = form({
+        name: 'foo',
+        children: [input({ name: 'baz' })],
+      });
+
+      const context = { foo: 'bar' };
+
+      const id = await rootMessenger.call(
+        'SnapInterfaceController:createInterface',
+        MOCK_SNAP_ID,
+        components,
+        context,
+      );
+
+      await rootMessenger.call(
+        'SnapInterfaceController:updateInterface',
+        MOCK_SNAP_ID,
+        id,
+        newContent,
+      );
+
+      const {
+        content,
+        state,
+        context: interfaceContext,
+      } = rootMessenger.call(
+        'SnapInterfaceController:getInterface',
+        MOCK_SNAP_ID,
+        id,
+      );
+
+      expect(content).toStrictEqual(getJsxElementFromComponent(newContent));
+      expect(state).toStrictEqual({ foo: { baz: null } });
+      expect(interfaceContext).toStrictEqual(context);
+    });
+
     it('throws if a link is on the phishing list', async () => {
       const rootMessenger = getRootSnapInterfaceControllerMessenger();
       const controllerMessenger = getRestrictedSnapInterfaceControllerMessenger(

packages/snaps-controllers/src/interface/SnapInterfaceController.ts

@@ -230,22 +230,28 @@ export class SnapInterfaceController extends BaseController<
    * @param snapId - The snap id requesting the update.
    * @param id - The interface id.
    * @param content - The new content.
+   * @param context - An optional interface context object.
    */
   async updateInterface(
     snapId: SnapId,
     id: string,
     content: ComponentOrElement,
+    context?: InterfaceContext,
   ) {
     this.#validateArgs(snapId, id);
     const element = getJsxInterface(content);
     await this.#validateContent(element);
+    validateInterfaceContext(context);
 
     const oldState = this.state.interfaces[id].state;
     const newState = constructState(oldState, element);
 
     this.update((draftState) => {
       draftState.interfaces[id].state = newState;
       draftState.interfaces[id].content = castDraft(element);
+      if (context) {
+        draftState.interfaces[id].context = context;
+      }
     });
   }
 

packages/snaps-rpc-methods/src/permitted/updateInterface.test.tsx

@@ -98,10 +98,58 @@ describe('snap_updateInterface', () => {
         <Box>
           <Text>Hello, world!</Text>
         </Box>,
+        undefined,
       );
     });
   });
 
+  it('updates the interface context', async () => {
+    const { implementation } = updateInterfaceHandler;
+
+    const updateInterface = jest.fn();
+
+    const hooks = {
+      updateInterface,
+    };
+
+    const engine = new JsonRpcEngine();
+
+    engine.push((request, response, next, end) => {
+      const result = implementation(
+        request as JsonRpcRequest<UpdateInterfaceParams>,
+        response as PendingJsonRpcResponse<UpdateInterfaceResult>,
+        next,
+        end,
+        hooks,
+      );
+
+      result?.catch(end);
+    });
+
+    await engine.handle({
+      jsonrpc: '2.0',
+      id: 1,
+      method: 'snap_updateInterface',
+      params: {
+        id: 'foo',
+        ui: (
+          <Box>
+            <Text>Hello, world!</Text>
+          </Box>
+        ) as JSXElement,
+        context: { foo: 'bar' },
+      },
+    });
+
+    expect(updateInterface).toHaveBeenCalledWith(
+      'foo',
+      <Box>
+        <Text>Hello, world!</Text>
+      </Box>,
+      { foo: 'bar' },
+    );
+  });
+
   it('throws on invalid params', async () => {
     const { implementation } = updateInterfaceHandler;
 

packages/snaps-rpc-methods/src/permitted/updateInterface.ts

@@ -6,10 +6,20 @@ import type {
   UpdateInterfaceResult,
   JsonRpcRequest,
   ComponentOrElement,
+  InterfaceContext,
+} from '@metamask/snaps-sdk';
+import {
+  ComponentOrElementStruct,
+  InterfaceContextStruct,
 } from '@metamask/snaps-sdk';
-import { ComponentOrElementStruct } from '@metamask/snaps-sdk';
 import { type InferMatching } from '@metamask/snaps-utils';
-import { StructError, create, object, string } from '@metamask/superstruct';
+import {
+  StructError,
+  create,
+  object,
+  optional,
+  string,
+} from '@metamask/superstruct';
 import type { PendingJsonRpcResponse } from '@metamask/utils';
 
 import type { MethodHooksObject } from '../utils';
@@ -22,8 +32,13 @@ export type UpdateInterfaceMethodHooks = {
   /**
    * @param id - The interface ID.
    * @param ui - The UI components.
+   * @param context - The optional interface context object.
    */
-  updateInterface: (id: string, ui: ComponentOrElement) => Promise<void>;
+  updateInterface: (
+    id: string,
+    ui: ComponentOrElement,
+    context?: InterfaceContext,
+  ) => Promise<void>;
 };
 
 export const updateInterfaceHandler: PermittedHandlerExport<
@@ -39,6 +54,7 @@ export const updateInterfaceHandler: PermittedHandlerExport<
 const UpdateInterfaceParametersStruct = object({
   id: string(),
   ui: ComponentOrElementStruct,
+  context: optional(InterfaceContextStruct),
 });
 
 export type UpdateInterfaceParameters = InferMatching<
@@ -70,9 +86,9 @@ async function getUpdateInterfaceImplementation(
   try {
     const validatedParams = getValidatedParams(params);
 
-    const { id, ui } = validatedParams;
+    const { id, ui, context } = validatedParams;
 
-    await updateInterface(id, ui);
+    await updateInterface(id, ui, context);
     res.result = null;
   } catch (error) {
     return end(error);

packages/snaps-sdk/src/types/methods/update-interface.ts

@@ -1,4 +1,4 @@
-import type { ComponentOrElement } from '..';
+import type { ComponentOrElement, InterfaceContext } from '..';
 
 /**
  * The request parameters for the `snap_createInterface` method.
@@ -9,6 +9,7 @@ import type { ComponentOrElement } from '..';
 export type UpdateInterfaceParams = {
   id: string;
   ui: ComponentOrElement;
+  context?: InterfaceContext;
 };
 
 /**