fix: shield coverage result polling (#6847)

## Explanation

<!--
Thanks for your contribution! Take a moment to answer these questions so
that reviewers have the information they need to properly understand
your changes:

* What is the current state of things and why does it need to change?
* What is the solution your changes offer and how does it work?
* Are there any changes whose purpose might not obvious to those
unfamiliar with the domain?
* If your primary goal was to update one package but you found you had
to update another one along the way, why did you do so?
* If you had to upgrade a dependency, why did you do so?
-->

This PR includes ~
- Updated internal coverage result polling and log logic.
   - Added cancellation logic to the polling.
   - Updated implementation of timeout.
   - Cancel any pending requests before starting new polling or logging.
- Updated TransactionMeta comparison in
`TransactionController:stateChange` subscriber to avoid triggering
multiple check coverage result unnecessarily.
- Removed `Personal Sign` check from the check signature coverage
result.

## References

<!--
Are there any issues that this pull request is tied to?
Are there other links that reviewers should consult to understand these
changes better?
Are there client or consumer pull requests to adopt any breaking
changes?

For example:

* Fixes #12345
* Related to #67890
-->

## Checklist

- [x] I've updated the test suite for new or updated code as appropriate
- [x] I've updated documentation (JSDoc, Markdown, etc.) for new or
updated code as appropriate
- [x] I've communicated my changes to consumers by [updating changelogs
for packages I've
changed](https://github.com/MetaMask/core/tree/main/docs/contributing.md#updating-changelogs),
highlighting breaking changes as necessary
- [x] I've prepared draft pull requests for clients and consumer
packages to resolve any breaking changes


<!-- CURSOR_SUMMARY -->
---

> [!NOTE]
> Replaces manual shield coverage polling with a
Cockatiel/controller-utils policy that supports retries and
cancellation, updates error handling, tests, changelog, and adds
required deps.
> 
> - **Backend**:
> - **Polling refactor**: Replace manual timeout loop in
`ShieldRemoteBackend.#getCoverageResult` with
`PollingWithCockatielPolicy`, using per-request IDs
(`txMeta.id`/`signatureRequest.id`).
> - **Cancellation**: Abort pending polls before
`logSignature`/`logTransaction` via
`#pollingPolicy.abortPendingRequest`.
> - **Retry/backoff**: Compute retries/backoff with `ConstantBackoff`
and `DEFAULT_MAX_RETRIES` (`computePollingIntervalAndRetryCount`).
> - **Error handling**: Throw `HttpError` with parsed body message for
non-200 responses.
> - **New Utility**:
> - Add `src/polling-with-policy.ts` implementing Cockatiel-based policy
with retry filtering (retry 4xx except 400) and request tracking.
> - **Tests**:
> - Update `backend.test.ts` expectations for new timeout/error
messages.
> - Add `polling-with-policy.test.ts` covering retry, max-retry,
cancellation, and concurrent requests.
>   - Add `tests/utils.delay` helper.
> - **Dependencies**:
>   - Add `@metamask/controller-utils` and `cockatiel`.
> - **Docs**:
>   - Update `CHANGELOG.md` noting optimized coverage-result polling.
> 
> <sup>Written by [Cursor
Bugbot](https://cursor.com/dashboard?tab=bugbot) for commit
e13e024. This will update automatically
on new commits. Configure
[here](https://cursor.com/dashboard?tab=bugbot).</sup>
<!-- /CURSOR_SUMMARY -->

Author: lwin-kyaw

packages/shield-controller/CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Fixed
+
+- Fixed and optimized shield-coverage-result polling with Cockatiel Policy from Controller-utils. ([#6847](https://github.com/MetaMask/core/pull/6847))
+
 ## [1.0.0]
 
 ### Added

packages/shield-controller/package.json

@@ -48,8 +48,10 @@
   },
   "dependencies": {
     "@metamask/base-controller": "^9.0.0",
+    "@metamask/controller-utils": "^11.14.1",
     "@metamask/messenger": "^0.3.0",
-    "@metamask/utils": "^11.8.1"
+    "@metamask/utils": "^11.8.1",
+    "cockatiel": "^3.1.2"
   },
   "devDependencies": {
     "@babel/runtime": "^7.23.9",

packages/shield-controller/src/backend.test.ts

@@ -129,7 +129,7 @@ describe('ShieldRemoteBackend', () => {
     expect(getAccessToken).toHaveBeenCalledTimes(1);
   });
 
-  it('should throw on check coverage timeout', async () => {
+  it('should throw on check coverage timeout with coverage status', async () => {
     const { backend, fetchMock } = setup({
       getCoverageResultTimeout: 0,
       getCoverageResultPollInterval: 0,
@@ -144,12 +144,42 @@ describe('ShieldRemoteBackend', () => {
     // Mock get coverage result: result unavailable.
     fetchMock.mockResolvedValue({
       status: 404,
-      json: jest.fn().mockResolvedValue({ status: 'unavailable' }),
     } as unknown as Response);
 
     const txMeta = generateMockTxMeta();
     await expect(backend.checkCoverage({ txMeta })).rejects.toThrow(
-      'Timeout waiting for coverage result',
+      'Failed to get coverage result: 404',
+    );
+
+    // Waiting here ensures coverage of the unexpected error and lets us know
+    // that the polling loop is exited as expected.
+    await new Promise((resolve) => setTimeout(resolve, 10));
+  });
+
+  it('should throw on check coverage timeout', async () => {
+    const { backend, fetchMock } = setup({
+      getCoverageResultTimeout: 0,
+      getCoverageResultPollInterval: 0,
+    });
+
+    // Mock init coverage check.
+    fetchMock.mockResolvedValueOnce({
+      status: 200,
+      json: jest.fn().mockResolvedValue({ coverageId: 'coverageId' }),
+    } as unknown as Response);
+
+    // Mock get coverage result: result unavailable.
+    fetchMock.mockResolvedValue({
+      status: 412,
+      json: jest.fn().mockResolvedValue({
+        message: 'Results are not available yet',
+        statusCode: 412,
+      }),
+    } as unknown as Response);
+
+    const txMeta = generateMockTxMeta();
+    await expect(backend.checkCoverage({ txMeta })).rejects.toThrow(
+      'Failed to get coverage result: Results are not available yet',
     );
 
     // Waiting here ensures coverage of the unexpected error and lets us know

packages/shield-controller/src/backend.ts

@@ -1,3 +1,8 @@
+import {
+  ConstantBackoff,
+  DEFAULT_MAX_RETRIES,
+  HttpError,
+} from '@metamask/controller-utils';
 import {
   EthMethod,
   SignatureRequestType,
@@ -7,6 +12,7 @@ import type { TransactionMeta } from '@metamask/transaction-controller';
 import type { Json } from '@metamask/utils';
 
 import { SignTypedDataVersion } from './constants';
+import { PollingWithCockatielPolicy } from './polling-with-policy';
 import type {
   CheckCoverageRequest,
   CheckSignatureCoverageRequest,
@@ -56,14 +62,12 @@ export type GetCoverageResultResponse = {
 export class ShieldRemoteBackend implements ShieldBackend {
   readonly #getAccessToken: () => Promise<string>;
 
-  readonly #getCoverageResultTimeout: number;
-
-  readonly #getCoverageResultPollInterval: number;
-
   readonly #baseUrl: string;
 
   readonly #fetch: typeof globalThis.fetch;
 
+  readonly #pollingPolicy: PollingWithCockatielPolicy;
+
   constructor({
     getAccessToken,
     getCoverageResultTimeout = 5000, // milliseconds
@@ -78,10 +82,18 @@ export class ShieldRemoteBackend implements ShieldBackend {
     fetch: typeof globalThis.fetch;
   }) {
     this.#getAccessToken = getAccessToken;
-    this.#getCoverageResultTimeout = getCoverageResultTimeout;
-    this.#getCoverageResultPollInterval = getCoverageResultPollInterval;
     this.#baseUrl = baseUrl;
     this.#fetch = fetchFn;
+
+    const { backoff, maxRetries } = computePollingIntervalAndRetryCount(
+      getCoverageResultTimeout,
+      getCoverageResultPollInterval,
+    );
+
+    this.#pollingPolicy = new PollingWithCockatielPolicy({
+      backoff,
+      maxRetries,
+    });
   }
 
   async checkCoverage(req: CheckCoverageRequest): Promise<CoverageResult> {
@@ -95,9 +107,11 @@ export class ShieldRemoteBackend implements ShieldBackend {
     }
 
     const txCoverageResultUrl = `${this.#baseUrl}/v1/transaction/coverage/result`;
-    const coverageResult = await this.#getCoverageResult(coverageId, {
-      coverageResultUrl: txCoverageResultUrl,
-    });
+    const coverageResult = await this.#getCoverageResult(
+      req.txMeta.id,
+      coverageId,
+      txCoverageResultUrl,
+    );
     return {
       coverageId,
       message: coverageResult.message,
@@ -119,9 +133,11 @@ export class ShieldRemoteBackend implements ShieldBackend {
     }
 
     const signatureCoverageResultUrl = `${this.#baseUrl}/v1/signature/coverage/result`;
-    const coverageResult = await this.#getCoverageResult(coverageId, {
-      coverageResultUrl: signatureCoverageResultUrl,
-    });
+    const coverageResult = await this.#getCoverageResult(
+      req.signatureRequest.id,
+      coverageId,
+      signatureCoverageResultUrl,
+    );
     return {
       coverageId,
       message: coverageResult.message,
@@ -138,6 +154,9 @@ export class ShieldRemoteBackend implements ShieldBackend {
       ...initBody,
     };
 
+    // cancel the pending get coverage result request
+    this.#pollingPolicy.abortPendingRequest(req.signatureRequest.id);
+
     const res = await this.#fetch(
       `${this.#baseUrl}/v1/signature/coverage/log`,
       {
@@ -159,6 +178,9 @@ export class ShieldRemoteBackend implements ShieldBackend {
       ...initBody,
     };
 
+    // cancel the pending get coverage result request
+    this.#pollingPolicy.abortPendingRequest(req.txMeta.id);
+
     const res = await this.#fetch(
       `${this.#baseUrl}/v1/transaction/coverage/log`,
       {
@@ -188,51 +210,39 @@ export class ShieldRemoteBackend implements ShieldBackend {
   }
 
   async #getCoverageResult(
+    requestId: string,
     coverageId: string,
-    configs: {
-      coverageResultUrl: string;
-      timeout?: number;
-      pollInterval?: number;
-    },
+    coverageResultUrl: string,
   ): Promise<GetCoverageResultResponse> {
     const reqBody: GetCoverageResultRequest = {
       coverageId,
     };
 
-    const timeout = configs?.timeout ?? this.#getCoverageResultTimeout;
-    const pollInterval =
-      configs?.pollInterval ?? this.#getCoverageResultPollInterval;
-
     const headers = await this.#createHeaders();
-    return await new Promise((resolve, reject) => {
-      let timeoutReached = false;
-      setTimeout(() => {
-        timeoutReached = true;
-        reject(new Error('Timeout waiting for coverage result'));
-      }, timeout);
-
-      const poll = async (): Promise<GetCoverageResultResponse> => {
-        // The timeoutReached variable is modified in the timeout callback.
-        // eslint-disable-next-line no-unmodified-loop-condition
-        while (!timeoutReached) {
-          const startTime = Date.now();
-          const res = await this.#fetch(configs.coverageResultUrl, {
-            method: 'POST',
-            headers,
-            body: JSON.stringify(reqBody),
-          });
-          if (res.status === 200) {
-            return (await res.json()) as GetCoverageResultResponse;
-          }
-          await sleep(pollInterval - (Date.now() - startTime));
-        }
-        // The following line will not have an effect as the upper level promise
-        // will already be rejected by now.
-        throw new Error('unexpected error');
-      };
-
-      poll().then(resolve).catch(reject);
-    });
+
+    const getCoverageResultFn = async (signal: AbortSignal) => {
+      const res = await this.#fetch(coverageResultUrl, {
+        method: 'POST',
+        headers,
+        body: JSON.stringify(reqBody),
+        signal,
+      });
+      if (res.status === 200) {
+        return (await res.json()) as GetCoverageResultResponse;
+      }
+
+      // parse the error message from the response body
+      let errorMessage = 'Timeout waiting for coverage result';
+      try {
+        const errorJson = await res.json();
+        errorMessage = `Failed to get coverage result: ${errorJson.message || errorJson.status}`;
+      } catch {
+        errorMessage = `Failed to get coverage result: ${res.status}`;
+      }
+      throw new HttpError(res.status, errorMessage);
+    };
+
+    return this.#pollingPolicy.start(requestId, getCoverageResultFn);
   }
 
   async #createHeaders() {
@@ -244,16 +254,6 @@ export class ShieldRemoteBackend implements ShieldBackend {
   }
 }
 
-/**
- * Sleep for a specified amount of time.
- *
- * @param ms - The number of milliseconds to sleep.
- * @returns A promise that resolves after the specified amount of time.
- */
-async function sleep(ms: number) {
-  return new Promise((resolve) => setTimeout(resolve, ms));
-}
-
 /**
  * Make the body for the init coverage check request.
  *
@@ -324,3 +324,28 @@ export function parseSignatureRequestMethod(
 
   return signatureRequest.type;
 }
+
+/**
+ * Compute the polling interval and retry count for the Cockatiel policy based on the timeout and poll interval given.
+ *
+ * @param timeout - The timeout in milliseconds.
+ * @param pollInterval - The poll interval in milliseconds.
+ * @returns The polling interval and retry count.
+ */
+function computePollingIntervalAndRetryCount(
+  timeout: number,
+  pollInterval: number,
+) {
+  const backoff = new ConstantBackoff(pollInterval);
+  const computedMaxRetries = Math.floor(timeout / pollInterval) + 1;
+
+  const maxRetries =
+    isNaN(computedMaxRetries) || !isFinite(computedMaxRetries)
+      ? DEFAULT_MAX_RETRIES
+      : computedMaxRetries;
+
+  return {
+    backoff,
+    maxRetries,
+  };
+}