feat(rpc): implement RPC URL override functionality in various modules

Author: Ansonhkg

e2e/src/e2e.spec.ts

@@ -19,6 +19,14 @@ import {
 } from './helper/tests';
 import { init } from './init';
 
+const RPC_OVERRIDE = process.env['LIT_YELLOWSTONE_PRIVATE_RPC_URL'];
+if (RPC_OVERRIDE) {
+  console.log(
+    '🧪 E2E: Using RPC override (LIT_YELLOWSTONE_PRIVATE_RPC_URL):',
+    RPC_OVERRIDE
+  );
+}
+
 describe('all', () => {
   // Singleton baby
   let ctx: Awaited<ReturnType<typeof init>>;
@@ -161,3 +169,99 @@ describe('all', () => {
     });
   });
 });
+
+describe('rpc override', () => {
+  const ORIGINAL_RPC = process.env.LIT_YELLOWSTONE_PRIVATE_RPC_URL;
+  const TEST_RPC = 'https://yellowstone-override.example';
+
+  beforeAll(() => {
+    process.env.LIT_YELLOWSTONE_PRIVATE_RPC_URL = TEST_RPC;
+  });
+
+  afterAll(() => {
+    process.env.LIT_YELLOWSTONE_PRIVATE_RPC_URL = ORIGINAL_RPC;
+  });
+
+  it('applies env rpc override to module and client', async () => {
+    const networks = await import('@lit-protocol/networks');
+
+    // choose module by NETWORK env (same way init.ts does)
+    const network = process.env.NETWORK || 'naga-dev';
+    const importNameMap: Record<string, string> = {
+      'naga-dev': 'nagaDev',
+      'naga-test': 'nagaTest',
+      'naga-local': 'nagaLocal',
+      'naga-staging': 'nagaStaging',
+    };
+    const importName = importNameMap[network];
+    const baseModule: any = (networks as any)[importName];
+
+    // apply override
+    const mod =
+      typeof baseModule.withOverrides === 'function'
+        ? baseModule.withOverrides({ rpcUrl: TEST_RPC })
+        : baseModule;
+
+    // log for verification
+    // base vs effective (when override is supported)
+    const baseRpcUrl =
+      typeof baseModule.getRpcUrl === 'function' ? baseModule.getRpcUrl() : 'n/a';
+    const effRpcUrl =
+      typeof mod.getRpcUrl === 'function' ? mod.getRpcUrl() : 'n/a';
+    // eslint-disable-next-line no-console
+    console.log('[rpc-override] TEST_RPC:', TEST_RPC);
+    // eslint-disable-next-line no-console
+    console.log('[rpc-override] module rpc (base → effective):', baseRpcUrl, '→', effRpcUrl);
+    try {
+      const baseChain =
+        typeof baseModule.getChainConfig === 'function' ? baseModule.getChainConfig() : null;
+      const effChain =
+        typeof mod.getChainConfig === 'function' ? mod.getChainConfig() : null;
+      if (baseChain && effChain) {
+        // eslint-disable-next-line no-console
+        console.log(
+          '[rpc-override] module chain id/name (base → effective):',
+          `${baseChain.id}/${baseChain.name}`,
+          '→',
+          `${effChain.id}/${effChain.name}`
+        );
+        // eslint-disable-next-line no-console
+        console.log(
+          '[rpc-override] module rpcUrls.default.http (base → effective):',
+          baseChain.rpcUrls.default.http,
+          '→',
+          effChain.rpcUrls.default.http
+        );
+        // eslint-disable-next-line no-console
+        console.log(
+          '[rpc-override] module rpcUrls.public.http (base → effective):',
+          (baseChain.rpcUrls as any)['public']?.http,
+          '→',
+          (effChain.rpcUrls as any)['public']?.http
+        );
+      }
+    } catch {}
+
+    // module reflects override
+    expect(mod.getRpcUrl()).toBe(TEST_RPC);
+    const chain = mod.getChainConfig();
+    expect(chain.rpcUrls.default.http[0]).toBe(TEST_RPC);
+    expect((chain.rpcUrls as any)['public'].http[0]).toBe(TEST_RPC);
+
+    // client reflects override
+    const { createLitClient } = await import('@lit-protocol/lit-client');
+    const client = await createLitClient({ network: mod });
+    const cc = client.getChainConfig();
+
+    // eslint-disable-next-line no-console
+    console.log('[rpc-override] client rpcUrl:', cc.rpcUrl);
+    // eslint-disable-next-line no-console
+    console.log('[rpc-override] client viem rpcUrls.default:', cc.viemConfig.rpcUrls.default.http);
+    // eslint-disable-next-line no-console
+    console.log('[rpc-override] client viem rpcUrls.public:', (cc.viemConfig.rpcUrls as any)['public']?.http);
+
+    expect(cc.rpcUrl).toBe(TEST_RPC);
+    expect(cc.viemConfig.rpcUrls.default.http[0]).toBe(TEST_RPC);
+    expect((cc.viemConfig.rpcUrls as any)['public'].http[0]).toBe(TEST_RPC);
+  });
+});

e2e/src/init.ts

@@ -109,7 +109,21 @@ export const init = async (
 
   // Dynamic import of network module
   const networksModule = await import('@lit-protocol/networks');
-  const _networkModule = networksModule[config.importName];
+  const _baseNetworkModule = networksModule[config.importName];
+
+  // Optional RPC override from env
+  const rpcOverride = process.env['LIT_YELLOWSTONE_PRIVATE_RPC_URL'];
+  const _networkModule =
+    rpcOverride && typeof _baseNetworkModule.withOverrides === 'function'
+      ? _baseNetworkModule.withOverrides({ rpcUrl: rpcOverride })
+      : _baseNetworkModule;
+
+  if (rpcOverride) {
+    console.log(
+      '✅ Using RPC override (LIT_YELLOWSTONE_PRIVATE_RPC_URL):',
+      rpcOverride
+    );
+  }
 
   // Fund accounts based on network type
   const isLocal = config.type === 'local';

packages/auth-services/src/_setup/initSystemContext.ts

@@ -15,7 +15,13 @@ declare global {
   };
 }
 
-export async function initSystemContext({ appName }: { appName: string }) {
+export async function initSystemContext({
+  appName,
+  rpcUrl,
+}: {
+  appName: string;
+  rpcUrl?: string;
+}) {
   console.log('🔥 [initSystemContext] Initializing system context...');
 
   let networkModule: any;
@@ -34,8 +40,14 @@ export async function initSystemContext({ appName }: { appName: string }) {
     throw new Error(`Unsupported network: ${env.NETWORK}`);
   }
 
+  // Apply runtime override if rpcUrl provided
+  const effectiveModule =
+    rpcUrl && typeof networkModule.withOverrides === 'function'
+      ? networkModule.withOverrides({ rpcUrl })
+      : networkModule;
+
   const litClient = await createLitClient({
-    network: networkModule,
+    network: effectiveModule,
   });
 
   const account = privateKeyToAccount(env.LIT_TXSENDER_PRIVATE_KEY as Hex);

packages/auth-services/src/auth-server/src/createAuthServer.ts

@@ -91,7 +91,10 @@ export const createLitAuthServer = (
       // =============================================================
       // Ensure system context is initialized before server fully starts, using appName from config
       // This was originally at the top level, moved here to use config.appName
-      await initSystemContext({ appName: config.appName });
+      await initSystemContext({
+        appName: config.appName,
+        rpcUrl: config.litTxsenderRpcUrl,
+      });
     })
     .get('/', () => ({
       message: 'PKP Auth Service is running. PKP minting is now asynchronous.',

packages/lit-client/src/lib/LitClient/createLitClient.ts

@@ -10,9 +10,9 @@ import {
 } from '@lit-protocol/access-control-conditions';
 import { encrypt as blsEncrypt } from '@lit-protocol/crypto';
 import { getChildLogger } from '@lit-protocol/logger';
-import type {
-  LitNetworkModule,
-  PKPStorageProvider,
+import {
+  type LitNetworkModule,
+  type PKPStorageProvider,
 } from '@lit-protocol/networks';
 import {
   AuthContextSchema2,

packages/lit-client/src/lib/LitClient/intergrations/createPkpViemAccount.ts

@@ -120,16 +120,14 @@ export async function createPKPViemAccount({
     tx,
     address,
     chain,
-    transportUrl,
   }: {
     tx: Partial<TransactionSerializable>;
     address: `0x${string}`;
     chain: Chain;
-    transportUrl: string;
   }): Promise<TransactionSerializable> {
     const client = createPublicClient({
       chain,
-      transport: http(transportUrl),
+      transport: http(chain.rpcUrls.default.http[0]),
     });
 
     try {
@@ -246,7 +244,6 @@ export async function createPKPViemAccount({
           tx: txRequest,
           address,
           chain: chainConfig,
-          transportUrl: chainConfig.rpcUrls.default.http[0],
         });
       } else {
         // Ensure minimum required fields for transaction type inference

packages/networks/src/networks/vNaga/shared/factories/BaseModuleFactory.ts

@@ -1087,6 +1087,72 @@ export function createBaseModule<T, M>(config: BaseModuleConfig<T, M>) {
         },
       },
     },
+    /**
+     * Returns a wrapped module instance with runtime overrides while keeping the base immutable.
+     * Currently supports overriding the RPC URL used by consumers of this module.
+     *
+     * @param overrides - The overrides to apply to the module.
+     * @returns A wrapped module instance with the overrides applied.
+     * @example
+     *
+     * import { nagaDev } from '@lit-protocol/networks';
+     * const nagaDevWithOverride = nagaDev.withOverrides({ rpcUrl: 'https://custom-rpc-url.com' });
+     * const litClient = await createLitClient({ network: nagaDevWithOverride });
+     */
+    withOverrides: (overrides: { rpcUrl?: string }) => {
+      const resolvedRpcUrl = overrides.rpcUrl ?? baseModule.getRpcUrl();
+
+      const wrapped = {
+        ...baseModule,
+
+        getRpcUrl: () => resolvedRpcUrl,
+
+        getChainConfig: () => {
+          const chain = baseModule.getChainConfig();
+          return {
+            ...chain,
+            rpcUrls: {
+              ...chain.rpcUrls,
+              default: {
+                ...chain.rpcUrls.default,
+                http: [resolvedRpcUrl],
+              },
+              ['public']: {
+                ...(chain.rpcUrls as any)['public'],
+                http: [resolvedRpcUrl],
+              },
+            },
+          } as typeof chain;
+        },
+
+        createStateManager: async <StateT, ModuleT>(params: {
+          callback: (params: CallbackParams) => Promise<StateT>;
+          networkModule: ModuleT;
+        }): Promise<Awaited<ReturnType<typeof createStateManager<StateT>>>> => {
+          const createReadOnlyChainManager = () => {
+            const dummyAccount = privateKeyToAccount(DEV_PRIVATE_KEY);
+            return createChainManager(dummyAccount);
+          };
+
+          const overriddenNetworkConfig = {
+            ...networkConfig,
+            rpcUrl: resolvedRpcUrl,
+            chainConfig: wrapped.getChainConfig(),
+          } as typeof networkConfig;
+
+          return await createStateManager<StateT>({
+            networkConfig: overriddenNetworkConfig,
+            callback: params.callback,
+            networkModule: wrapped as unknown as LitNetworkModuleBase,
+            createReadOnlyChainManager,
+          });
+        },
+      } as typeof baseModule & {
+        withOverrides: (o: { rpcUrl?: string }) => any;
+      };
+
+      return wrapped;
+    },
   };
 
   return baseModule;