feat: add deleteAfterResult option to approval controller (#4715)

packages/approval-controller/src/ApprovalController.test.ts

@@ -4,6 +4,7 @@ import { ControllerMessenger } from '@metamask/base-controller';
 import { errorCodes, JsonRpcError } from '@metamask/rpc-errors';
 import { nanoid } from 'nanoid';
 
+import { flushPromises } from '../../../tests/helpers';
 import type {
  AddApprovalOptions,
  ApprovalControllerActions,
@@ -975,6 +976,93 @@ describe('approval controller', () => {
  approvalController.has({ origin: 'bar.baz' }),
  ).toBe(true);
  });
+
+ describe('with deleteAfterResult', () => {
+ it('deletes entry after result callback', async () => {
+ const approvalPromise = approvalController.add({
+ id: ID_MOCK,
+ origin: ORIGIN_MOCK,
+ type: TYPE,
+ expectsResult: true,
+ });
+
+ const resultPromise = approvalController.accept(ID_MOCK, VALUE_MOCK, {
+ waitForResult: true,
+ deleteAfterResult: true,
+ });
+
+ const { resultCallbacks } = await approvalPromise;
+
+ await flushPromises();
+
+ expect(approvalController.has({ id: ID_MOCK })).toBe(true);
+
+ resultCallbacks?.success(RESULT_MOCK);
+
+ await flushPromises();
+
+ expect(approvalController.has({ id: ID_MOCK })).toBe(false);
+
+ await resultPromise;
+ });
+
+ it.each([false, undefined])(
+ 'deletes immediately if waitForResult is %s',
+ async (waitForResult) => {
+ const approvalPromise = approvalController.add({
+ id: ID_MOCK,
+ origin: ORIGIN_MOCK,
+ type: TYPE,
+ expectsResult: true,
+ });
+
+ const resultPromise = approvalController.accept(ID_MOCK, VALUE_MOCK, {
+ waitForResult,
+ deleteAfterResult: true,
+ });
+
+ const { resultCallbacks } = await approvalPromise;
+
+ await flushPromises();
+
+ expect(approvalController.has({ id: ID_MOCK })).toBe(false);
+
+ resultCallbacks?.success(RESULT_MOCK);
+
+ await flushPromises();
+
+ expect(approvalController.has({ id: ID_MOCK })).toBe(false);
+
+ await resultPromise;
+ },
+ );
+
+ it('throws if already accepted', async () => {
+ const approvalPromise = approvalController.add({
+ id: ID_MOCK,
+ origin: ORIGIN_MOCK,
+ type: TYPE,
+ expectsResult: true,
+ });
+
+ const resultPromise = approvalController.accept(ID_MOCK, VALUE_MOCK, {
+ waitForResult: true,
+ deleteAfterResult: true,
+ });
+
+ const { resultCallbacks } = await approvalPromise;
+
+ await flushPromises();
+
+ await approvalController.accept(ID_MOCK, VALUE_MOCK);
+
+ resultCallbacks?.success(RESULT_MOCK);
+
+ await expect(resultPromise).rejects.toThrow(
+ `Approval request with id '${ID_MOCK}' not found.`,
+ );
+ });
+ });
  });
 
  describe('reject', () => {

packages/approval-controller/src/ApprovalController.ts

@@ -181,6 +181,13 @@ export type AcceptOptions = {
  * If false or unspecified, the promise will resolve immediately.
  */
  waitForResult?: boolean;
+
+ /**
+ * Whether to delete the approval request after a result callback is called.
+ * If false or unspecified, the approval request will be deleted immediately.
+ * Ignored if `waitForResult` is false or unspecified.
+ */
+ deleteAfterResult?: boolean;
 };
 
 export type StartFlowOptions = OptionalField<
@@ -688,9 +695,15 @@ export class ApprovalController extends BaseController<
  ): Promise<AcceptResult> {
  // Safe to cast as the delete method below will throw if the ID is not found
  const approval = this.get(id) as ApprovalRequest<ApprovalRequestData>;
- const requestPromise = this.#deleteApprovalAndGetCallbacks(id);
+ const requestPromise = this.#getCallbacks(id);
+ let requestDeleted = false;
 
- return new Promise((resolve, reject) => {
+ if (!options?.deleteAfterResult || !options.waitForResult) {
+ this.#delete(id);
+ requestDeleted = true;
+ }
+
+ return new Promise<AcceptResult>((resolve, reject) => {
  const resultCallbacks: AcceptResultCallbacks = {
  success: (acceptValue?: unknown) => resolve({ value: acceptValue }),
  error: reject,
@@ -712,6 +725,10 @@ export class ApprovalController extends BaseController<
  if (!options?.waitForResult) {
  resolve({ value: undefined });
  }
+ }).finally(() => {
+ if (!requestDeleted) {
+ this.#delete(id);
+ }
  });
  }
 
@@ -723,7 +740,9 @@ export class ApprovalController extends BaseController<
  * @param error - The error to reject the approval promise with.
  */
  reject(id: string, error: unknown): void {
- this.#deleteApprovalAndGetCallbacks(id).reject(error);
+ const callbacks = this.#getCallbacks(id);
+ this.#delete(id);
+ callbacks.reject(error);
  }
 
  /**
@@ -1022,20 +1041,21 @@ export class ApprovalController extends BaseController<
  }
 
  /**
- * Deletes the approval with the given id. The approval promise must be
- * resolved or reject before this method is called.
+ * Deletes the approval with the given id.
+ *
  * Deletion is an internal operation because approval state is solely
  * managed by this controller.
  *
  * @param id - The id of the approval request to be deleted.
  */
  #delete(id: string): void {
+ if (!this.#approvals.has(id)) {
+ throw new ApprovalRequestNotFoundError(id);
+ }
+
  this.#approvals.delete(id);
 
- // This method is only called after verifying that the approval with the
- // specified id exists.
- // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
- const { origin, type } = this.state.pendingApprovals[id]!;
+ const { origin, type } = this.state.pendingApprovals[id];
 
  const originMap = this.#origins.get(origin) as Map<string, number>;
  const originTotalCount = this.getApprovalCount({ origin });
@@ -1055,21 +1075,13 @@ export class ApprovalController extends BaseController<
  });
  }
 
- /**
- * Gets the approval callbacks for the given id, deletes the entry, and then
- * returns the callbacks for promise resolution.
- * Throws an error if no approval is found for the given id.
- *
- * @param id - The id of the approval request.
- * @returns The promise callbacks associated with the approval request.
- */
- #deleteApprovalAndGetCallbacks(id: string): ApprovalCallbacks {
+ #getCallbacks(id: string): ApprovalCallbacks {
  const callbacks = this.#approvals.get(id);
+
  if (!callbacks) {
  throw new ApprovalRequestNotFoundError(id);
  }
 
- this.#delete(id);
  return callbacks;
  }