From c83bba15568b488b14a1cdd3cb4b3016c3d7c4bf Mon Sep 17 00:00:00 2001
From: Callum <callum.mcintyre@anza.xyz>
Date: Fri, 6 Sep 2024 09:56:18 +0100
Subject: [PATCH] Allow extending numeric keypaths when extending Solana RPC
 API

---
 .changeset/tricky-flies-brush.md              |  5 ++
 .../__tests__/create-solana-rpc-api.test.ts   | 48 +++++++++++++++++++
 .../create-solana-rpc-api-type-test.ts.ts     | 46 ++++++++++++++++++
 packages/rpc-api/src/index.ts                 | 15 ++++--
 packages/rpc/src/rpc-default-config.ts        |  4 +-
 5 files changed, 113 insertions(+), 5 deletions(-)
 create mode 100644 .changeset/tricky-flies-brush.md
 create mode 100644 packages/rpc-api/src/__tests__/create-solana-rpc-api.test.ts
 create mode 100644 packages/rpc-api/src/__typetests__/create-solana-rpc-api-type-test.ts.ts

diff --git a/.changeset/tricky-flies-brush.md b/.changeset/tricky-flies-brush.md
new file mode 100644
index 000000000000..41217b77511c
--- /dev/null
+++ b/.changeset/tricky-flies-brush.md
@@ -0,0 +1,5 @@
+---
+'@solana/rpc-api': patch
+---
+
+Allow extending numeric keypaths when extending Solana RPC API
diff --git a/packages/rpc-api/src/__tests__/create-solana-rpc-api.test.ts b/packages/rpc-api/src/__tests__/create-solana-rpc-api.test.ts
new file mode 100644
index 000000000000..0c6f84a417c1
--- /dev/null
+++ b/packages/rpc-api/src/__tests__/create-solana-rpc-api.test.ts
@@ -0,0 +1,48 @@
+import { getDefaultResponseTransformerForSolanaRpc } from '@solana/rpc-transformers';
+
+import { createSolanaRpcApi, SolanaRpcApi } from '..';
+
+jest.mock('@solana/rpc-transformers');
+
+describe('createSolanaRpcApi', () => {
+    describe('when extending allowedNumericKeyPaths', () => {
+        it('should create a response transformer including the new paths', () => {
+            type TestApi = SolanaRpcApi & {
+                someNewFunction(): void;
+                someOtherFunction(): void;
+            };
+
+            createSolanaRpcApi<TestApi>({
+                extendAllowedNumericKeypaths: {
+                    someNewFunction: [['someNewKeyPath']],
+                    someOtherFunction: [['someOtherKeyPath']],
+                },
+            });
+
+            expect(getDefaultResponseTransformerForSolanaRpc as jest.Mock).toHaveBeenCalledWith({
+                allowedNumericKeyPaths: expect.objectContaining({
+                    someNewFunction: [['someNewKeyPath']],
+                    someOtherFunction: [['someOtherKeyPath']],
+                }),
+            });
+        });
+
+        it('should include the default paths', () => {
+            type TestApi = SolanaRpcApi & {
+                someNewFunction(): void;
+            };
+
+            createSolanaRpcApi<TestApi>({
+                extendAllowedNumericKeypaths: {
+                    someNewFunction: [['someNewKeyPath']],
+                },
+            });
+
+            expect(getDefaultResponseTransformerForSolanaRpc as jest.Mock).toHaveBeenCalledWith({
+                allowedNumericKeyPaths: expect.objectContaining({
+                    getAccountInfo: expect.any(Array),
+                }),
+            });
+        });
+    });
+});
diff --git a/packages/rpc-api/src/__typetests__/create-solana-rpc-api-type-test.ts.ts b/packages/rpc-api/src/__typetests__/create-solana-rpc-api-type-test.ts.ts
new file mode 100644
index 000000000000..b8fecdac8e12
--- /dev/null
+++ b/packages/rpc-api/src/__typetests__/create-solana-rpc-api-type-test.ts.ts
@@ -0,0 +1,46 @@
+import { createSolanaRpcApi, SolanaRpcApi } from '..';
+
+type TestApi = SolanaRpcApi & {
+    someNewFunction(): void;
+    someOtherFunction(): void;
+};
+
+createSolanaRpcApi<TestApi>();
+createSolanaRpcApi<TestApi>({
+    extendAllowedNumericKeypaths: {},
+});
+createSolanaRpcApi<TestApi>({
+    extendAllowedNumericKeypaths: {
+        someNewFunction: [['someNewKeyPath']],
+    },
+});
+createSolanaRpcApi<TestApi>({
+    extendAllowedNumericKeypaths: {
+        someNewFunction: [['someNewKeyPath']],
+        someOtherFunction: [['someOtherKeyPath']],
+    },
+});
+
+createSolanaRpcApi<TestApi>({
+    extendAllowedNumericKeypaths: {
+        // @ts-expect-error getAccountInfo is a SolanaRpcApi method
+        getAccountInfo: [],
+    },
+});
+
+createSolanaRpcApi<SolanaRpcApi>();
+createSolanaRpcApi<SolanaRpcApi>({
+    extendAllowedNumericKeypaths: {},
+});
+createSolanaRpcApi<SolanaRpcApi>({
+    extendAllowedNumericKeypaths: {
+        // @ts-expect-error getAccountInfo is a SolanaRpcApi method
+        getAccountInfo: [],
+    },
+});
+createSolanaRpcApi<SolanaRpcApi>({
+    extendAllowedNumericKeypaths: {
+        // @ts-expect-error someNewMethod is not a method on the RPC API
+        someNewMethod: [['someNewKeyPath']],
+    },
+});
diff --git a/packages/rpc-api/src/index.ts b/packages/rpc-api/src/index.ts
index 7423f99d7f52..e811f807229e 100644
--- a/packages/rpc-api/src/index.ts
+++ b/packages/rpc-api/src/index.ts
@@ -179,15 +179,24 @@ export type {
     SimulateTransactionApi,
 };
 
-type Config = RequestTransformerConfig;
+// This is required because `AllowedKeyPaths<never>` allows any method
+// So we instead force it to be an empty object in that case
+type StrictAllowedNumericKeypaths<T> = [keyof T] extends [never] ? Record<string, never> : AllowedNumericKeypaths<T>;
+
+type Config<
+    TRpcMethods extends SolanaRpcApi | SolanaRpcApiDevnet | SolanaRpcApiMainnet | SolanaRpcApiTestnet = SolanaRpcApi,
+> = Readonly<{
+    extendAllowedNumericKeypaths?: StrictAllowedNumericKeypaths<Omit<TRpcMethods, keyof SolanaRpcApi>>;
+}> &
+    RequestTransformerConfig;
 
 export function createSolanaRpcApi<
     TRpcMethods extends SolanaRpcApi | SolanaRpcApiDevnet | SolanaRpcApiMainnet | SolanaRpcApiTestnet = SolanaRpcApi,
->(config?: Config): RpcApi<TRpcMethods> {
+>(config?: Config<TRpcMethods>): RpcApi<TRpcMethods> {
     return createRpcApi<TRpcMethods>({
         requestTransformer: getDefaultRequestTransformerForSolanaRpc(config),
         responseTransformer: getDefaultResponseTransformerForSolanaRpc({
-            allowedNumericKeyPaths: getAllowedNumericKeypaths(),
+            allowedNumericKeyPaths: { ...(config?.extendAllowedNumericKeypaths ?? {}), ...getAllowedNumericKeypaths() },
         }),
     });
 }
diff --git a/packages/rpc/src/rpc-default-config.ts b/packages/rpc/src/rpc-default-config.ts
index efd42cfa03e2..9dba1ce6aedd 100644
--- a/packages/rpc/src/rpc-default-config.ts
+++ b/packages/rpc/src/rpc-default-config.ts
@@ -2,9 +2,9 @@ import type { createSolanaRpcApi } from '@solana/rpc-api';
 
 import { createSolanaJsonRpcIntegerOverflowError } from './rpc-integer-overflow-error';
 
-export const DEFAULT_RPC_CONFIG: Partial<NonNullable<Parameters<typeof createSolanaRpcApi>[0]>> = {
+export const DEFAULT_RPC_CONFIG = {
     defaultCommitment: 'confirmed',
     onIntegerOverflow(request, keyPath, value) {
         throw createSolanaJsonRpcIntegerOverflowError(request.methodName, keyPath, value);
     },
-};
+} as const satisfies Partial<NonNullable<Parameters<typeof createSolanaRpcApi>[0]>>;