feat: Add context parameter to InternalProvider and JsonRpcServer methods

package.json

@@ -63,6 +63,7 @@
     "@metamask/eth-block-tracker": "^14.0.0",
     "@metamask/eth-json-rpc-provider": "^5.0.1",
     "@metamask/json-rpc-engine": "^10.1.1",
+    "@metamask/network-controller": "^25.0.0",
     "@metamask/utils": "^11.8.1",
     "@ts-bridge/cli": "^0.6.4",
     "@types/jest": "^27.4.1",

packages/bridge-controller/src/utils/balance.test.ts

@@ -2,16 +2,16 @@ import { BigNumber } from '@ethersproject/bignumber';
 import { AddressZero } from '@ethersproject/constants';
 import { Contract } from '@ethersproject/contracts';
 import { Web3Provider } from '@ethersproject/providers';
-import type { InternalProvider } from '@metamask/eth-json-rpc-provider';
 import { abiERC20 } from '@metamask/metamask-eth-abis';
+import type { Provider } from '@metamask/network-controller';
 
 import * as balanceUtils from './balance';
 import { fetchTokenBalance } from './balance';
 import { FakeProvider } from '../../../../tests/fake-provider';
 
 declare global {
   // eslint-disable-next-line no-var
-  var ethereumProvider: InternalProvider;
+  var ethereumProvider: Provider;
 }
 
 jest.mock('@ethersproject/contracts', () => {

packages/eth-block-tracker/CHANGELOG.md

@@ -7,6 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- Add `Context` generic parameter to `PollingBlockTracker` ([#7061](https://github.com/MetaMask/core/pull/7061))
+  - This enables passing providers with different context types to the block tracker.
+
 ### Changed
 
 - **BREAKING:** Use `InternalProvider` instead of `SafeEventEmitterProvider` ([#6796](https://github.com/MetaMask/core/pull/6796))

packages/eth-block-tracker/src/PollingBlockTracker.ts

@@ -1,4 +1,8 @@
 import type { InternalProvider } from '@metamask/eth-json-rpc-provider';
+import type {
+  ContextConstraint,
+  MiddlewareContext,
+} from '@metamask/json-rpc-engine/v2';
 import SafeEventEmitter from '@metamask/safe-event-emitter';
 import {
   createDeferredPromise,
@@ -17,8 +21,10 @@ const sec = 1000;
 
 const blockTrackerEvents: (string | symbol)[] = ['sync', 'latest'];
 
-export type PollingBlockTrackerOptions = {
-  provider?: InternalProvider;
+export type PollingBlockTrackerOptions<
+  Context extends ContextConstraint = MiddlewareContext,
+> = {
+  provider?: InternalProvider<Context>;
   pollingInterval?: number;
   retryTimeout?: number;
   keepEventLoopActive?: boolean;
@@ -33,7 +39,9 @@ type ExtendedJsonRpcRequest = {
 
 type InternalListener = (value: string) => void;
 
-export class PollingBlockTracker
+export class PollingBlockTracker<
+    Context extends ContextConstraint = MiddlewareContext,
+  >
   extends SafeEventEmitter
   implements BlockTracker
 {
@@ -49,7 +57,7 @@ export class PollingBlockTracker
 
   private _pollingTimeout?: ReturnType<typeof setTimeout>;
 
-  private readonly _provider: InternalProvider;
+  private readonly _provider: InternalProvider<Context>;
 
   private readonly _pollingInterval: number;
 
@@ -65,7 +73,7 @@ export class PollingBlockTracker
 
   #pendingFetch?: Omit<DeferredPromise<string>, 'resolve'>;
 
-  constructor(opts: PollingBlockTrackerOptions = {}) {
+  constructor(opts: PollingBlockTrackerOptions<Context> = {}) {
     // parse + validate args
     if (!opts.provider) {
       throw new Error('PollingBlockTracker - no provider specified.');

packages/eth-json-rpc-provider/CHANGELOG.md

@@ -11,6 +11,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Add `providerFromMiddlewareV2` ([#7001](https://github.com/MetaMask/core/pull/7001))
   - This accepts the new middleware from `@metamask/json-rpc-engine/v2`.
+- Add `context` option to `InternalProvider.request()` ([#7061](https://github.com/MetaMask/core/pull/7061))
+  - Enables passing a `MiddlewareContext` to the JSON-RPC server.
 
 ### Changed
 

packages/eth-json-rpc-provider/src/internal-provider.test.ts

@@ -2,7 +2,10 @@ import { Web3Provider } from '@ethersproject/providers';
 import EthQuery from '@metamask/eth-query';
 import EthJsQuery from '@metamask/ethjs-query';
 import { asV2Middleware, JsonRpcEngine } from '@metamask/json-rpc-engine';
-import type { JsonRpcMiddleware } from '@metamask/json-rpc-engine/v2';
+import type {
+  JsonRpcMiddleware,
+  MiddlewareContext,
+} from '@metamask/json-rpc-engine/v2';
 import { JsonRpcEngineV2 } from '@metamask/json-rpc-engine/v2';
 import { providerErrors, rpcErrors } from '@metamask/rpc-errors';
 import { type JsonRpcRequest, type Json } from '@metamask/utils';
@@ -16,7 +19,9 @@ import {
 
 jest.mock('uuid');
 
-type ResultParam = Json | ((req?: JsonRpcRequest) => Json);
+type ResultParam =
+  | Json
+  | ((req?: JsonRpcRequest, context?: MiddlewareContext) => Json);
 
 const createLegacyEngine = (method: string, result: ResultParam) => {
   const engine = new JsonRpcEngine();
@@ -33,9 +38,11 @@ const createLegacyEngine = (method: string, result: ResultParam) => {
 const createV2Engine = (method: string, result: ResultParam) => {
   return JsonRpcEngineV2.create<JsonRpcMiddleware<JsonRpcRequest>>({
     middleware: [
-      ({ request, next }) => {
+      ({ request, next, context }) => {
         if (request.method === method) {
-          return typeof result === 'function' ? result(request) : result;
+          return typeof result === 'function'
+            ? result(request as JsonRpcRequest, context)
+            : result;
         }
         return next();
       },
@@ -245,6 +252,29 @@ describe.each([
       expect(response.result).toBe(42);
     });
 
+    it('forwards the context to the JSON-RPC handler', async () => {
+      const rpcHandler = createRpcHandler('test', (request, context) => {
+        // @ts-expect-error - Intentional type abuse.
+        // eslint-disable-next-line jest/no-conditional-in-test
+        return context?.assertGet('foo') ?? request.foo;
+      });
+      const provider = new InternalProvider({ engine: rpcHandler });
+
+      const request = {
+        id: 1,
+        jsonrpc: '2.0' as const,
+        method: 'test',
+      };
+
+      const result = await provider.request(request, {
+        context: {
+          foo: 'bar',
+        },
+      });
+
+      expect(result).toBe('bar');
+    });
+
     it('handles a successful EIP-1193 object request', async () => {
       let req: JsonRpcRequest | undefined;
       const rpcHandler = createRpcHandler('test', (request) => {

packages/eth-json-rpc-provider/src/internal-provider.ts

@@ -1,9 +1,10 @@
 import { asV2Middleware, type JsonRpcEngine } from '@metamask/json-rpc-engine';
-import type {
-  ContextConstraint,
-  MiddlewareContext,
+import {
+  type HandleOptions,
+  type ContextConstraint,
+  type MiddlewareContext,
+  JsonRpcEngineV2,
 } from '@metamask/json-rpc-engine/v2';
-import { JsonRpcEngineV2 } from '@metamask/json-rpc-engine/v2';
 import type { JsonRpcSuccess } from '@metamask/utils';
 import {
   type Json,
@@ -37,16 +38,18 @@ type Options<
  * This provider loosely follows conventions that pre-date EIP-1193.
  * It is not compliant with any Ethereum provider standard.
  */
-export class InternalProvider {
-  readonly #engine: JsonRpcEngineV2<JsonRpcRequest, MiddlewareContext>;
+export class InternalProvider<
+  Context extends ContextConstraint = MiddlewareContext,
+> {
+  readonly #engine: JsonRpcEngineV2<JsonRpcRequest, Context>;
 
   /**
    * Construct a InternalProvider from a JSON-RPC server or legacy engine.
    *
    * @param options - Options.
    * @param options.engine - The JSON-RPC engine used to process requests.
    */
-  constructor({ engine }: Options) {
+  constructor({ engine }: Options<JsonRpcRequest, Context>) {
     this.#engine =
       'push' in engine
         ? JsonRpcEngineV2.create({
@@ -59,14 +62,17 @@ export class InternalProvider {
    * Send a provider request asynchronously.
    *
    * @param eip1193Request - The request to send.
+   * @param options - The options for the request operation.
+   * @param options.context - The context to include with the request.
    * @returns The JSON-RPC response.
    */
   async request<Params extends JsonRpcParams, Result extends Json>(
     eip1193Request: Eip1193Request<Params>,
+    options?: HandleOptions<Context>,
   ): Promise<Result> {
     const jsonRpcRequest =
       convertEip1193RequestToJsonRpcRequest(eip1193Request);
-    return (await this.#handle<Result>(jsonRpcRequest)).result;
+    return (await this.#handle<Result>(jsonRpcRequest, options)).result;
   }
 
   /**
@@ -118,12 +124,14 @@ export class InternalProvider {
 
   readonly #handle = async <Result extends Json>(
     jsonRpcRequest: JsonRpcRequest,
+    options?: HandleOptions<Context>,
   ): Promise<JsonRpcSuccess<Result>> => {
     const { id, jsonrpc } = jsonRpcRequest;
-    // This typecast is technicaly unsafe, but we need it to preserve the provider's
-    // public interface, which allows you to typecast results.
+    // The `result` typecast is unsafe, but we need it to preserve the provider's
+    // public interface, which allows you to unsafely typecast results.
     const result = (await this.#engine.handle(
       jsonRpcRequest,
+      options,
     )) as unknown as Result;
 
     return {

packages/json-rpc-engine/CHANGELOG.md

@@ -9,7 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
-- Add `JsonRpcEngineV2` ([#6176](https://github.com/MetaMask/core/pull/6176), [#6971](https://github.com/MetaMask/core/pull/6971), [#6975](https://github.com/MetaMask/core/pull/6975), [#6990](https://github.com/MetaMask/core/pull/6990), [#6991](https://github.com/MetaMask/core/pull/6991), [#7001](https://github.com/MetaMask/core/pull/7001), [#7032](https://github.com/MetaMask/core/pull/7032))
+- Add `JsonRpcEngineV2` ([#6176](https://github.com/MetaMask/core/pull/6176), [#6971](https://github.com/MetaMask/core/pull/6971), [#6975](https://github.com/MetaMask/core/pull/6975), [#6990](https://github.com/MetaMask/core/pull/6990), [#6991](https://github.com/MetaMask/core/pull/6991), [#7032](https://github.com/MetaMask/core/pull/7032), [#7001](https://github.com/MetaMask/core/pull/7001), [#7061](https://github.com/MetaMask/core/pull/7061))
   - This is a complete rewrite of `JsonRpcEngine`, intended to replace the original implementation. See the readme for details.
 
 ## [10.1.1]

packages/json-rpc-engine/README.md

@@ -444,6 +444,33 @@ const engine = JsonRpcEngineV2.create({
 });
 ```
 
+#### Passing the context to `handle()`
+
+You can pass a `MiddlewareContext` instance directly to `handle()`:
+
+```ts
+const context = new MiddlewareContext();
+context.set('foo', 'bar');
+const result = await engine.handle(
+  { id: '1', jsonrpc: '2.0', method: 'hello' },
+  { context },
+);
+console.log(result); // 'bar'
+```
+
+You can also pass a plain object as a shorthand for a `MiddlewareContext` instance:
+
+```ts
+const context = { foo: 'bar' };
+const result = await engine.handle(
+  { id: '1', jsonrpc: '2.0', method: 'hello' },
+  { context },
+);
+console.log(result); // 'bar'
+```
+
+This works the same way for `JsonRpcServer.handle()`.
+
 #### Constraining context keys and values
 
 The context exposes a generic parameter `KeyValues`, which determines the keys and values

packages/json-rpc-engine/src/v2/JsonRpcEngineV2.test.ts

@@ -371,6 +371,24 @@ describe('JsonRpcEngineV2', () => {
         expect(result).toBe('bar');
       });
 
+      it('accepts an initial context as a KeyValues object', async () => {
+        const initialContext = { foo: 'bar' } as const;
+        const middleware: JsonRpcMiddleware<
+          JsonRpcRequest,
+          string,
+          MiddlewareContext<Record<string, string>>
+        > = ({ context }) => context.assertGet('foo');
+        const engine = JsonRpcEngineV2.create({
+          middleware: [middleware],
+        });
+
+        const result = await engine.handle(makeRequest(), {
+          context: initialContext,
+        });
+
+        expect(result).toBe('bar');
+      });
+
       it('accepts middleware with different context types', async () => {
         const middleware1: JsonRpcMiddleware<
           JsonRpcCall,

packages/json-rpc-engine/src/v2/JsonRpcEngineV2.ts

@@ -7,7 +7,11 @@ import {
 } from '@metamask/utils';
 import deepFreeze from 'deep-freeze-strict';
 
-import type { ContextConstraint, MergeContexts } from './MiddlewareContext';
+import type {
+  ContextConstraint,
+  InferKeyValues,
+  MergeContexts,
+} from './MiddlewareContext';
 import { MiddlewareContext } from './MiddlewareContext';
 import {
   isNotification,
@@ -60,8 +64,11 @@ type RequestState<Request extends JsonRpcCall> = {
   result: Readonly<ResultConstraint<Request>> | undefined;
 };
 
-type HandleOptions<Context extends MiddlewareContext> = {
-  context?: Context;
+/**
+ * The options for the JSON-RPC request/notification handling operation.
+ */
+export type HandleOptions<Context extends ContextConstraint> = {
+  context?: Context | InferKeyValues<Context>;
 };
 
 type ConstructorOptions<
@@ -286,12 +293,14 @@ export class JsonRpcEngineV2<
    * operation. Permits returning an `undefined` result.
    *
    * @param originalRequest - The JSON-RPC request to handle.
-   * @param context - The context to pass to the middleware.
+   * @param rawContext - The context to pass to the middleware.
    * @returns The result from the middleware.
    */
   async #handle(
     originalRequest: Request,
-    context: Context = new MiddlewareContext() as Context,
+    rawContext:
+      | Context
+      | InferKeyValues<Context> = new MiddlewareContext() as Context,
   ): Promise<RequestState<Request>> {
     this.#assertIsNotDestroyed();
 
@@ -303,6 +312,9 @@ export class JsonRpcEngineV2<
     };
     const middlewareIterator = this.#makeMiddlewareIterator();
     const firstMiddleware = middlewareIterator.next().value;
+    const context = MiddlewareContext.isInstance(rawContext)
+      ? rawContext
+      : (new MiddlewareContext(rawContext) as Context);
 
     const makeNext = this.#makeNextFactory(middlewareIterator, state, context);