feat: Add snap_startTrace and snap_endTrace methods for performance tracing

packages/examples/packages/preinstalled/snap.manifest.json

@@ -7,7 +7,7 @@
     "url": "https://github.com/MetaMask/snaps.git"
   },
   "source": {
-    "shasum": "orIGELPnJuNGSXC/IH7XrDciP6y+csrN9oWh//+m7DI=",
+    "shasum": "yNJdBPGXVZJZHIjQtrlPGXA+BlDBZhjZIswZbTjuYkM=",
     "location": {
       "npm": {
         "filePath": "dist/bundle.js",

packages/examples/packages/preinstalled/src/index.tsx

@@ -64,6 +64,23 @@ export const onRpcRequest: OnRpcRequestHandler = async ({ request }) => {
       });
     }
 
+    case 'startTrace':
+      return await snap.request({
+        method: 'snap_startTrace',
+        params: {
+          name: 'Test Snap Trace',
+        },
+      });
+
+    case 'endTrace': {
+      return await snap.request({
+        method: 'snap_endTrace',
+        params: {
+          name: 'Test Snap Trace',
+        },
+      });
+    }
+
     default:
       throw new MethodNotFoundError({ method: request.method });
   }

packages/snaps-rpc-methods/jest.config.js

@@ -10,10 +10,10 @@ module.exports = deepmerge(baseConfig, {
   ],
   coverageThreshold: {
     global: {
-      branches: 95.32,
-      functions: 98.73,
-      lines: 98.89,
-      statements: 98.59,
+      branches: 95.37,
+      functions: 98.76,
+      lines: 98.92,
+      statements: 98.62,
     },
   },
 });

packages/snaps-rpc-methods/src/permitted/endTrace.test.ts

@@ -0,0 +1,173 @@
+import { JsonRpcEngine } from '@metamask/json-rpc-engine';
+import type { EndTraceParams, EndTraceResult } from '@metamask/snaps-sdk';
+import type { JsonRpcRequest, PendingJsonRpcResponse } from '@metamask/utils';
+
+import { endTraceHandler } from './endTrace';
+
+describe('snap_endTrace', () => {
+  describe('endTraceHandler', () => {
+    it('has the expected shape', () => {
+      expect(endTraceHandler).toMatchObject({
+        methodNames: ['snap_endTrace'],
+        implementation: expect.any(Function),
+        hookNames: {
+          endTrace: true,
+          getSnap: true,
+        },
+      });
+    });
+  });
+
+  describe('implementation', () => {
+    it('calls the `endTrace` hook with the provided parameters', async () => {
+      const { implementation } = endTraceHandler;
+
+      const endTrace = jest.fn().mockReturnValue(null);
+
+      const getSnap = jest.fn().mockReturnValue({ preinstalled: true });
+      const hooks = { endTrace, getSnap };
+
+      const engine = new JsonRpcEngine();
+
+      engine.push((request, response, next, end) => {
+        const result = implementation(
+          request as JsonRpcRequest<EndTraceParams>,
+          response as PendingJsonRpcResponse<EndTraceResult>,
+          next,
+          end,
+          hooks,
+        );
+
+        result?.catch(end);
+      });
+
+      const response = await engine.handle({
+        jsonrpc: '2.0',
+        id: 1,
+        method: 'snap_endTrace',
+        params: {
+          id: 'test-id',
+          name: 'Test Trace',
+          timestamp: 1234567890,
+        },
+      });
+
+      expect(response).toStrictEqual({
+        jsonrpc: '2.0',
+        id: 1,
+        result: null,
+      });
+
+      expect(endTrace).toHaveBeenCalledWith({
+        id: 'test-id',
+        name: 'Test Trace',
+        timestamp: 1234567890,
+      });
+    });
+
+    it('throws an error if the Snap is not preinstalled', async () => {
+      const { implementation } = endTraceHandler;
+
+      const endTrace = jest.fn();
+      const getSnap = jest.fn().mockReturnValue({ preinstalled: false });
+      const hooks = { endTrace, getSnap };
+
+      const engine = new JsonRpcEngine();
+
+      engine.push((request, response, next, end) => {
+        const result = implementation(
+          request as JsonRpcRequest<EndTraceParams>,
+          response as PendingJsonRpcResponse<EndTraceResult>,
+          next,
+          end,
+          hooks,
+        );
+
+        result?.catch(end);
+      });
+
+      const response = await engine.handle({
+        jsonrpc: '2.0',
+        id: 1,
+        method: 'snap_endTrace',
+        params: {
+          id: 'test-id',
+          name: 'Test Trace',
+        },
+      });
+
+      expect(endTrace).not.toHaveBeenCalled();
+      expect(response).toStrictEqual({
+        jsonrpc: '2.0',
+        id: 1,
+        error: {
+          code: -32601,
+          message: 'The method does not exist / is not available.',
+          stack: expect.any(String),
+        },
+      });
+    });
+
+    it.each([
+      [
+        { foo: 'bar' },
+        'Invalid params: At path: name -- Expected a string, but received: undefined.',
+      ],
+      [
+        { name: undefined },
+        'Invalid params: At path: name -- Expected a string, but received: undefined.',
+      ],
+      [
+        { name: 'Test Trace', id: 123 },
+        'Invalid params: At path: id -- Expected a string, but received: 123.',
+      ],
+      [
+        { name: 'Test Trace', id: 'test-id', timestamp: 'not-a-number' },
+        'Invalid params: At path: timestamp -- Expected a number, but received: "not-a-number".',
+      ],
+    ])(
+      'throws an error if the parameters are invalid',
+      async (params, error) => {
+        const { implementation } = endTraceHandler;
+
+        const endTrace = jest.fn();
+        const getSnap = jest.fn().mockReturnValue({ preinstalled: true });
+        const hooks = { endTrace, getSnap };
+
+        const engine = new JsonRpcEngine();
+
+        engine.push((request, response, next, end) => {
+          const result = implementation(
+            request as JsonRpcRequest<EndTraceParams>,
+            response as PendingJsonRpcResponse<EndTraceResult>,
+            next,
+            end,
+            hooks,
+          );
+
+          result?.catch(end);
+        });
+
+        // @ts-expect-error: Intentionally passing invalid params.
+        // eslint-disable-next-line @typescript-eslint/await-thenable
+        const response = await engine.handle({
+          jsonrpc: '2.0',
+          id: 1,
+          method: 'snap_endTrace',
+          params,
+        });
+
+        expect(endTrace).not.toHaveBeenCalled();
+        expect(response).toStrictEqual({
+          jsonrpc: '2.0',
+          id: 1,
+          error: {
+            code: -32602,
+            message: error,
+            stack: expect.any(String),
+          },
+        });
+      },
+    );
+  });
+});

packages/snaps-rpc-methods/src/permitted/endTrace.ts

@@ -0,0 +1,132 @@
+import type { JsonRpcEngineEndCallback } from '@metamask/json-rpc-engine';
+import type { PermittedHandlerExport } from '@metamask/permission-controller';
+import { rpcErrors } from '@metamask/rpc-errors';
+import type {
+  JsonRpcRequest,
+  EndTraceParams,
+  EndTraceResult,
+  EndTraceRequest,
+} from '@metamask/snaps-sdk';
+import type { InferMatching, Snap } from '@metamask/snaps-utils';
+import {
+  number,
+  create,
+  object,
+  string,
+  StructError,
+  exactOptional,
+} from '@metamask/superstruct';
+import type { PendingJsonRpcResponse } from '@metamask/utils';
+
+import type { MethodHooksObject } from '../utils';
+
+const hookNames: MethodHooksObject<EndTraceMethodHooks> = {
+  endTrace: true,
+  getSnap: true,
+};
+
+export type EndTraceMethodHooks = {
+  /**
+   * End a performance trace in Sentry.
+   *
+   * @param request - The trace request object.
+   * @returns The performance trace context.
+   */
+  endTrace: (request: EndTraceRequest) => void;
+
+  /**
+   * Get Snap metadata.
+   *
+   * @param snapId - The ID of a Snap.
+   */
+  getSnap: (snapId: string) => Snap | undefined;
+};
+
+const EndTraceParametersStruct = object({
+  id: exactOptional(string()),
+  name: string(),
+  timestamp: exactOptional(number()),
+});
+
+export type EndTraceParameters = InferMatching<
+  typeof EndTraceParametersStruct,
+  EndTraceParams
+>;
+
+/**
+ * Handler for the `snap_endTrace` method.
+ */
+export const endTraceHandler: PermittedHandlerExport<
+  EndTraceMethodHooks,
+  EndTraceParameters,
+  EndTraceResult
+> = {
+  methodNames: ['snap_endTrace'],
+  implementation: getEndTraceImplementation,
+  hookNames,
+};
+
+/**
+ * The `snap_endTrace` method implementation. This method is used to end a
+ * performance trace in Sentry. It is only available to preinstalled Snaps.
+ *
+ * @param request - The JSON-RPC request object.
+ * @param response - The JSON-RPC response object.
+ * @param _next - The `json-rpc-engine` "next" callback. Not used by this
+ * function.
+ * @param end - The `json-rpc-engine` "end" callback.
+ * @param hooks - The RPC method hooks.
+ * @param hooks.endTrace - The hook function to end a performance trace.
+ * @param hooks.getSnap - The hook function to get Snap metadata.
+ * @returns Nothing.
+ */
+function getEndTraceImplementation(
+  request: JsonRpcRequest<EndTraceParameters>,
+  response: PendingJsonRpcResponse,
+  _next: unknown,
+  end: JsonRpcEngineEndCallback,
+  { endTrace, getSnap }: EndTraceMethodHooks,
+): void {
+  const snap = getSnap(
+    (request as JsonRpcRequest<EndTraceParams> & { origin: string }).origin,
+  );
+
+  if (!snap?.preinstalled) {
+    return end(rpcErrors.methodNotFound());
+  }
+
+  const { params } = request;
+
+  try {
+    const validatedParams = getValidatedParams(params);
+    endTrace(validatedParams);
+
+    response.result = null;
+  } catch (error) {
+    return end(error);
+  }
+
+  return end();
+}
+
+/**
+ * Validate the parameters for the `snap_endTrace` method.
+ *
+ * @param params - Parameters to validate.
+ * @returns Validated parameters.
+ * @throws Throws RPC error if validation fails.
+ */
+function getValidatedParams(params: unknown): EndTraceParameters {
+  try {
+    return create(params, EndTraceParametersStruct);
+  } catch (error) {
+    if (error instanceof StructError) {
+      throw rpcErrors.invalidParams({
+        message: `Invalid params: ${error.message}.`,
+      });
+    }
+
+    /* istanbul ignore next */
+    throw rpcErrors.internal();
+  }
+}

packages/snaps-rpc-methods/src/permitted/handlers.ts

@@ -2,6 +2,7 @@ import { cancelBackgroundEventHandler } from './cancelBackgroundEvent';
 import { clearStateHandler } from './clearState';
 import { closeWebSocketHandler } from './closeWebSocket';
 import { createInterfaceHandler } from './createInterface';
+import { endTraceHandler } from './endTrace';
 import { providerRequestHandler } from './experimentalProviderRequest';
 import { getAllSnapsHandler } from './getAllSnaps';
 import { getBackgroundEventsHandler } from './getBackgroundEvents';
@@ -22,6 +23,7 @@ import { resolveInterfaceHandler } from './resolveInterface';
 import { scheduleBackgroundEventHandler } from './scheduleBackgroundEvent';
 import { sendWebSocketMessageHandler } from './sendWebSocketMessage';
 import { setStateHandler } from './setState';
+import { startTraceHandler } from './startTrace';
 import { trackErrorHandler } from './trackError';
 import { trackEventHandler } from './trackEvent';
 import { updateInterfaceHandler } from './updateInterface';
@@ -55,6 +57,8 @@ export const methodHandlers = {
   snap_closeWebSocket: closeWebSocketHandler,
   snap_sendWebSocketMessage: sendWebSocketMessageHandler,
   snap_getWebSockets: getWebSocketsHandler,
+  snap_startTrace: startTraceHandler,
+  snap_endTrace: endTraceHandler,
 };
 /* eslint-enable @typescript-eslint/naming-convention */