MetaMask · MajorLift · Oct 5, 2023
diff --git a/packages/eth-json-rpc-provider/src/index.ts b/packages/eth-json-rpc-provider/src/index.ts
@@ -1,3 +1,3 @@
 export * from './provider-from-engine';
 export * from './provider-from-middleware';
-export type { SafeEventEmitterProvider } from './safe-event-emitter-provider';
+export { SafeEventEmitterProvider } from './safe-event-emitter-provider';
diff --git a/packages/network-controller/package.json b/packages/network-controller/package.json
@@ -35,6 +35,7 @@
     "@metamask/eth-json-rpc-middleware": "^11.0.2",
     "@metamask/eth-json-rpc-provider": "^2.2.0",
     "@metamask/eth-query": "^3.0.1",
+    "@metamask/json-rpc-engine": "^7.1.1",
     "@metamask/swappable-obj-proxy": "^2.1.0",
     "@metamask/utils": "^8.1.0",
     "async-mutex": "^0.2.6",

diff --git a/packages/network-controller/src/NetworkController.ts b/packages/network-controller/src/NetworkController.ts
@@ -8,9 +8,12 @@ import {
   NetworkType,
   isSafeChainId,
 } from '@metamask/controller-utils';
-import EthQuery from '@metamask/eth-query';
+import EthQuery, { type Provider } from '@metamask/eth-query';
 import { createEventEmitterProxy } from '@metamask/swappable-obj-proxy';
-import type { SwappableProxy } from '@metamask/swappable-obj-proxy';
+import type {
+  EventEmitterLike,
+  SwappableProxy,
+} from '@metamask/swappable-obj-proxy';
 import type { Hex } from '@metamask/utils';
 import {
   assertIsStrictHexString,
@@ -32,7 +35,6 @@ import { projectLogger, createModuleLogger } from './logger';
 import { NetworkClientType } from './types';
 import type {
   BlockTracker,
-  Provider,
   CustomNetworkClientConfiguration,
   InfuraNetworkClientConfiguration,
   NetworkClientConfiguration,
@@ -357,7 +359,9 @@ export type BlockTrackerProxy = SwappableProxy<
  * selected network can change without consumers needing to refresh the object
  * reference to that network.)
  */
-export type ProviderProxy = SwappableProxy<ProxyWithAccessibleTarget<Provider>>;
+export type ProviderProxy =
+  | SwappableProxy<ProxyWithAccessibleTarget<Provider>>
+  | SwappableProxy<EventEmitterLike>;
 
 export type NetworkControllerStateChangeEvent = {
   type: `NetworkController:stateChange`;

diff --git a/packages/network-controller/src/create-auto-managed-network-client.ts b/packages/network-controller/src/create-auto-managed-network-client.ts
@@ -1,10 +1,12 @@
-import type { NetworkClient } from './create-network-client';
-import { createNetworkClient } from './create-network-client';
 import type {
-  BlockTracker,
-  NetworkClientConfiguration,
   Provider,
-} from './types';
+  ProviderSendAsyncCallback,
+  SendAsyncPayload,
+} from '@metamask/eth-query';
+
+import type { NetworkClient } from './create-network-client';
+import { createNetworkClient } from './create-network-client';
+import type { BlockTracker, NetworkClientConfiguration } from './types';
 
 /**
  * The name of the method on both the provider and block tracker proxy which can
@@ -90,10 +92,13 @@ export function createAutoManagedNetworkClient<
           // Ensure that the method on the provider is called with `this` as
           // the target, *not* the proxy (which happens by default) —
           // this allows private properties to be accessed
-          return function (this: unknown, ...args: any[]) {
-            // @ts-expect-error We don't care that `this` may not be compatible
-            // with the signature of the method being called, as technically
-            // it can be anything.
+          return function (
+            this: unknown,
+            ...args: [
+              payload: SendAsyncPayload<unknown>,
+              callback: ProviderSendAsyncCallback<unknown>,
+            ]
+          ) {
             return value.apply(this === receiver ? provider : this, args);
           };
         }

diff --git a/packages/network-controller/src/create-network-client.ts b/packages/network-controller/src/create-network-client.ts
@@ -10,26 +10,23 @@ import {
   createFetchMiddleware,
   createRetryOnEmptyMiddleware,
 } from '@metamask/eth-json-rpc-middleware';
-import type { SafeEventEmitterProvider } from '@metamask/eth-json-rpc-provider';
 import {
+  type SafeEventEmitterProvider,
   providerFromEngine,
   providerFromMiddleware,
 } from '@metamask/eth-json-rpc-provider';
-import type { Hex } from '@metamask/utils';
-import { PollingBlockTracker } from 'eth-block-tracker';
+import type { Provider } from '@metamask/eth-query';
 import {
   createAsyncMiddleware,
   createScaffoldMiddleware,
   JsonRpcEngine,
   mergeMiddleware,
-} from 'json-rpc-engine';
-import type { JsonRpcMiddleware } from 'json-rpc-engine';
-
-import type {
-  BlockTracker,
-  NetworkClientConfiguration,
-  Provider,
-} from './types';
+} from '@metamask/json-rpc-engine';
+import type { JsonRpcMiddleware } from '@metamask/json-rpc-engine';
+import type { Hex, Json, JsonRpcParams } from '@metamask/utils';
+import { PollingBlockTracker } from 'eth-block-tracker';
+
+import type { BlockTracker, NetworkClientConfiguration } from './types';
 import { NetworkClientType } from './types';
 
 const SECOND = 1000;
@@ -56,19 +53,21 @@ export function createNetworkClient(
 ): NetworkClient {
   const rpcApiMiddleware =
     networkConfig.type === NetworkClientType.Infura
-      ? createInfuraMiddleware({
+      ? createInfuraMiddleware<JsonRpcParams, Json>({
           network: networkConfig.network,
           projectId: networkConfig.infuraProjectId,
           maxAttempts: 5,
           source: 'metamask',
         })
-      : createFetchMiddleware({
+      : createFetchMiddleware<JsonRpcParams, Json>({
           btoa: global.btoa,
           fetch: global.fetch,
           rpcUrl: networkConfig.rpcUrl,
         });
 
-  const rpcProvider = providerFromMiddleware(rpcApiMiddleware);
+  const rpcProvider = providerFromMiddleware<JsonRpcParams, Json>(
+    rpcApiMiddleware,
+  );
 
   const blockTrackerOpts =
     // eslint-disable-next-line n/no-process-env
@@ -82,13 +81,13 @@ export function createNetworkClient(
 
   const networkMiddleware =
     networkConfig.type === NetworkClientType.Infura
-      ? createInfuraNetworkMiddleware({
+      ? createInfuraNetworkMiddleware<JsonRpcParams, Json>({
           blockTracker,
           network: networkConfig.network,
           rpcProvider,
           rpcApiMiddleware,
         })
-      : createCustomNetworkMiddleware({
+      : createCustomNetworkMiddleware<JsonRpcParams, Json>({
           blockTracker,
           chainId: networkConfig.chainId,
           rpcApiMiddleware,
@@ -117,7 +116,10 @@ export function createNetworkClient(
  * @param args.rpcApiMiddleware - Additional middleware.
  * @returns The collection of middleware that makes up the Infura client.
  */
-function createInfuraNetworkMiddleware({
+function createInfuraNetworkMiddleware<
+  Params extends JsonRpcParams | unknown = JsonRpcParams,
+  Result extends Json[] | Record<string, Json> | unknown = unknown,
+>({
   blockTracker,
   network,
   rpcProvider,
@@ -126,7 +128,7 @@ function createInfuraNetworkMiddleware({
   blockTracker: PollingBlockTracker;
   network: InfuraNetworkType;
   rpcProvider: SafeEventEmitterProvider;
-  rpcApiMiddleware: JsonRpcMiddleware<unknown, unknown>;
+  rpcApiMiddleware: JsonRpcMiddleware<Params, Result>;
 }) {
   return mergeMiddleware([
     createNetworkAndChainIdMiddleware({ network }),
@@ -156,9 +158,9 @@ function createNetworkAndChainIdMiddleware({
   });
 }
 
-const createChainIdMiddleware = (
+const createChainIdMiddleware = <Params = unknown>(
   chainId: Hex,
-): JsonRpcMiddleware<unknown, unknown> => {
+): JsonRpcMiddleware<Params, `0x${string}`> => {
   return (req, res, next, end) => {
     if (req.method === 'eth_chainId') {
       res.result = chainId;
@@ -177,14 +179,17 @@ const createChainIdMiddleware = (
  * @param args.rpcApiMiddleware - Additional middleware.
  * @returns The collection of middleware that makes up the Infura client.
  */
-function createCustomNetworkMiddleware({
+function createCustomNetworkMiddleware<
+  Params extends JsonRpcParams | unknown = unknown,
+  Result extends Json[] | Record<string, Json> | unknown = unknown,
+>({
   blockTracker,
   chainId,
   rpcApiMiddleware,
 }: {
   blockTracker: PollingBlockTracker;
   chainId: Hex;
-  rpcApiMiddleware: any;
+  rpcApiMiddleware: JsonRpcMiddleware<Params, Result>;
 }) {
   // eslint-disable-next-line n/no-process-env
   const testMiddlewares = process.env.IN_TEST

diff --git a/packages/network-controller/src/index.ts b/packages/network-controller/src/index.ts
@@ -1,5 +1,7 @@
 export * from './NetworkController';
 export * from './constants';
-export type { BlockTracker, Provider } from './types';
-export type { NetworkClientConfiguration } from './types';
-export { NetworkClientType } from './types';
+export type {
+  BlockTracker,
+  NetworkClientConfiguration,
+  NetworkClientType,
+} from './types';
diff --git a/packages/network-controller/src/types.ts b/packages/network-controller/src/types.ts
@@ -1,10 +1,7 @@
 import type { InfuraNetworkType } from '@metamask/controller-utils';
-import type { SafeEventEmitterProvider } from '@metamask/eth-json-rpc-provider';
 import type { Hex } from '@metamask/utils';
 import type { PollingBlockTracker } from 'eth-block-tracker';
 
-export type Provider = SafeEventEmitterProvider;
-
 export type BlockTracker = PollingBlockTracker;
 
 /**

diff --git a/packages/network-controller/tests/NetworkController.test.ts b/packages/network-controller/tests/NetworkController.test.ts
@@ -7,6 +7,7 @@ import {
   NetworkType,
   toHex,
 } from '@metamask/controller-utils';
+import type { Provider } from '@metamask/eth-query';
 import assert from 'assert';
 import { ethErrors } from 'eth-rpc-errors';
 import type { Patch } from 'immer';
@@ -27,10 +28,8 @@ import type {
   ProviderConfig,
 } from '../src/NetworkController';
 import { NetworkController } from '../src/NetworkController';
-import type { Provider } from '../src/types';
 import { NetworkClientType } from '../src/types';
-import type { FakeProviderStub } from './fake-provider';
-import { FakeProvider } from './fake-provider';
+import { FakeProvider, type FakeProviderStub } from './fake-provider';
 
 jest.mock('../src/create-network-client');
 
@@ -978,9 +977,10 @@ describe('NetworkController', () => {
                 provider,
               );
               const response1 = await promisifiedSendAsync1({
-                id: '1',
+                id: 1,
                 jsonrpc: '2.0',
                 method: 'test',
+                params: [],
               });
               expect(response1.result).toBe('test response 1');
 
@@ -989,9 +989,10 @@ describe('NetworkController', () => {
                 provider,
               );
               const response2 = await promisifiedSendAsync2({
-                id: '2',
+                id: 2,
                 jsonrpc: '2.0',
                 method: 'test',
+                params: [],
               });
               expect(response2.result).toBe('test response 2');
             },
@@ -1072,9 +1073,10 @@ describe('NetworkController', () => {
               provider,
             );
             const response1 = await promisifiedSendAsync1({
-              id: '1',
+              id: 1,
               jsonrpc: '2.0',
               method: 'test',
+              params: [],
             });
             expect(response1.result).toBe('test response 1');
 
@@ -1083,9 +1085,10 @@ describe('NetworkController', () => {
               provider,
             );
             const response2 = await promisifiedSendAsync2({
-              id: '2',
+              id: 2,
               jsonrpc: '2.0',
               method: 'test',
+              params: [],
             });
             expect(response2.result).toBe('test response 2');
           },
@@ -4948,9 +4951,10 @@ describe('NetworkController', () => {
                   provider,
                 );
                 const response = await promisifiedSendAsync({
-                  id: '1',
+                  id: 1,
                   jsonrpc: '2.0',
                   method: 'test',
+                  params: [],
                 });
                 expect(response.result).toBe('test response');
               },
@@ -5431,10 +5435,10 @@ describe('NetworkController', () => {
                 // We only care about the first state change, because it
                 // happens before networkDidChange
                 count: 1,
-                operation: () => {
+                operation: async () => {
                   // Intentionally not awaited because we want to check state
                   // while this operation is in-progress
-                  controller.rollbackToPreviousProvider();
+                  await controller.rollbackToPreviousProvider();
                 },
                 beforeResolving: () => {
                   expect(
@@ -5502,9 +5506,10 @@ describe('NetworkController', () => {
                 provider,
               );
               const response = await promisifiedSendAsync({
-                id: '1',
+                id: 1,
                 jsonrpc: '2.0',
                 method: 'test',
+                params: [],
               });
               expect(response.result).toBe('test response');
             },