From f63c3d3ed50f9a73374af26f1a93879b0501fb34 Mon Sep 17 00:00:00 2001
From: Elliot Winkler <elliot.winkler@gmail.com>
Date: Wed, 22 Jan 2025 17:18:48 -0700
Subject: [PATCH] Add conveniences to createServicePolicy

Expose these exports from `cockatiel` so consumers can use them without
needing to depend on the library directly:

- `BrokenCircuitError`
- `CircuitState`
- `handleAll`
- `handleWhen`

Also, update `createServicePolicy` to expose the retry and circuit
breaker policies so that we can access their state in other modules that
use this utility function.

This is needed to be able to create an RPC service.
---
 packages/controller-utils/CHANGELOG.md        |  3 ++-
 packages/controller-utils/jest.config.js      |  2 +-
 .../src/create-service-policy.ts              | 19 +++++++++++++++++--
 packages/controller-utils/src/index.test.ts   |  4 ++++
 packages/controller-utils/src/index.ts        |  4 ++++
 5 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/packages/controller-utils/CHANGELOG.md b/packages/controller-utils/CHANGELOG.md
index ef8364c1754..725f8b50829 100644
--- a/packages/controller-utils/CHANGELOG.md
+++ b/packages/controller-utils/CHANGELOG.md
@@ -9,10 +9,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
-- Add utility function for reducing boilerplate for service classes ([#5154](https://github.com/MetaMask/core/pull/5154), [#5143](https://github.com/MetaMask/core/pull/5143), [#5149](https://github.com/MetaMask/core/pull/5149), [#5188](https://github.com/MetaMask/core/pull/5188))
+- Add utility function for reducing boilerplate for service classes ([#5154](https://github.com/MetaMask/core/pull/5154), [#5143](https://github.com/MetaMask/core/pull/5143), [#5149](https://github.com/MetaMask/core/pull/5149), [#5188](https://github.com/MetaMask/core/pull/5188), [#5192](https://github.com/MetaMask/core/pull/5192))
   - Add function `createServicePolicy`
   - Add constants `DEFAULT_CIRCUIT_BREAK_DURATION`, `DEFAULT_DEGRADED_THRESHOLD`, `DEFAULT_MAX_CONSECUTIVE_FAILURES`, and `DEFAULT_MAX_RETRIES`
   - Add types `ServicePolicy` and `CreateServicePolicyOptions`
+  - Re-export `BrokenCircuitError`, `CircuitState`, `handleAll`, and `handleWhen` from `cockatiel`
 
 ## [11.4.5]
 
diff --git a/packages/controller-utils/jest.config.js b/packages/controller-utils/jest.config.js
index 746ae98e6f5..926cf376355 100644
--- a/packages/controller-utils/jest.config.js
+++ b/packages/controller-utils/jest.config.js
@@ -18,7 +18,7 @@ module.exports = merge(baseConfig, {
   coverageThreshold: {
     global: {
       branches: 78.12,
-      functions: 84.61,
+      functions: 80.7,
       lines: 87.3,
       statements: 86.5,
     },
diff --git a/packages/controller-utils/src/create-service-policy.ts b/packages/controller-utils/src/create-service-policy.ts
index 971fae8f8cb..2747418d480 100644
--- a/packages/controller-utils/src/create-service-policy.ts
+++ b/packages/controller-utils/src/create-service-policy.ts
@@ -1,11 +1,13 @@
 import {
-  circuitBreaker,
+  BrokenCircuitError,
+  CircuitState,
   ConsecutiveBreaker,
   ExponentialBackoff,
+  circuitBreaker,
   handleAll,
+  handleWhen,
   retry,
   wrap,
-  CircuitState,
 } from 'cockatiel';
 import type {
   CircuitBreakerPolicy,
@@ -14,6 +16,8 @@ import type {
   RetryPolicy,
 } from 'cockatiel';
 
+export { CircuitState, BrokenCircuitError, handleAll, handleWhen };
+
 /**
  * The options for `createServicePolicy`.
  */
@@ -51,6 +55,15 @@ export type CreateServicePolicyOptions = {
  * The service policy object.
  */
 export type ServicePolicy = IPolicy & {
+  /**
+   * The Cockatiel circuit breaker policy that the service policy uses
+   * internally.
+   */
+  circuitBreakerPolicy: CircuitBreakerPolicy;
+  /**
+   * The Cockatiel retry policy that the service policy uses internally.
+   */
+  retryPolicy: RetryPolicy;
   /**
    * A function which is called when the number of times that the service fails
    * in a row meets the set maximum number of consecutive failures.
@@ -218,6 +231,8 @@ export function createServicePolicy({
 
   return {
     ...policy,
+    circuitBreakerPolicy,
+    retryPolicy,
     onBreak,
     onDegraded,
     onRetry,
diff --git a/packages/controller-utils/src/index.test.ts b/packages/controller-utils/src/index.test.ts
index 5573baf5125..f6db054b3ce 100644
--- a/packages/controller-utils/src/index.test.ts
+++ b/packages/controller-utils/src/index.test.ts
@@ -4,11 +4,15 @@ describe('@metamask/controller-utils', () => {
   it('has expected JavaScript exports', () => {
     expect(Object.keys(allExports)).toMatchInlineSnapshot(`
       Array [
+        "BrokenCircuitError",
+        "CircuitState",
         "DEFAULT_CIRCUIT_BREAK_DURATION",
         "DEFAULT_DEGRADED_THRESHOLD",
         "DEFAULT_MAX_CONSECUTIVE_FAILURES",
         "DEFAULT_MAX_RETRIES",
         "createServicePolicy",
+        "handleAll",
+        "handleWhen",
         "BNToHex",
         "convertHexToDecimal",
         "fetchWithErrorHandling",
diff --git a/packages/controller-utils/src/index.ts b/packages/controller-utils/src/index.ts
index 67742049deb..f0f65a9b1da 100644
--- a/packages/controller-utils/src/index.ts
+++ b/packages/controller-utils/src/index.ts
@@ -1,9 +1,13 @@
 export {
+  BrokenCircuitError,
+  CircuitState,
   DEFAULT_CIRCUIT_BREAK_DURATION,
   DEFAULT_DEGRADED_THRESHOLD,
   DEFAULT_MAX_CONSECUTIVE_FAILURES,
   DEFAULT_MAX_RETRIES,
   createServicePolicy,
+  handleAll,
+  handleWhen,
 } from './create-service-policy';
 export type {
   CreateServicePolicyOptions,