fix(private-rpc): withOverrides no longer monkey-patches methods nor wraps every chain API. It now builds an overridden networkConfig with the resolved RPC URL and a cloned chainConfig whose rpcUrls.default.http and rpcUrls['public'].http use the override. All downstream consumers (state manager, chain APIs, contracts) automatically use the overridden RPC via the standard construction path. No per-method wrapping needed.

Ansonhkg · Ansonhkg · commit 18cc2c2f812d · 2025-09-05T14:20:30.000+01:00
diff --git a/packages/networks/src/chains/ChronicleYellowstone.ts b/packages/networks/src/chains/ChronicleYellowstone.ts
@@ -41,3 +41,27 @@ export const WagmiConfig = createConfig({
     [viemChainConfig.id]: http(),
   },
 });
+
+/**
+ * Resolve the effective RPC URL from an optional override.
+ */
+export function resolveRpcUrl(overrideRpc?: string): string {
+  return overrideRpc ?? RPC_URL;
+}
+
+/**
+ * Build a Chain config using the base Chronicle Yellowstone config,
+ * applying an optional RPC override to both default and public http URLs.
+ */
+export function buildViemChainConfig(overrideRpc?: string): Chain {
+  const rpc = resolveRpcUrl(overrideRpc);
+  const base = viemChainConfig;
+  return {
+    ...base,
+    rpcUrls: {
+      ...base.rpcUrls,
+      default: { ...base.rpcUrls.default, http: [rpc] },
+      public: { ...(base.rpcUrls as any)['public'], http: [rpc] },
+    },
+  } as Chain;
+}
diff --git a/packages/networks/src/networks/vNaga/envs/base/BaseNetworkEnvironment.ts b/packages/networks/src/networks/vNaga/envs/base/BaseNetworkEnvironment.ts
@@ -15,6 +15,7 @@ export interface BaseEnvironmentOptions<T, M> {
   minimumThreshold?: number;
   httpProtocol?: 'http://' | 'https://';
   requiredAttestation?: boolean;
+  rpcUrlOverride?: string;
 }
 
 export abstract class BaseNetworkEnvironment<T, M> {
@@ -24,9 +25,9 @@ export abstract class BaseNetworkEnvironment<T, M> {
     this.config = {
       minimumThreshold: options.minimumThreshold || 3,
       network: options.network,
-      rpcUrl: this.getRpcUrl(),
+      rpcUrl: this.getRpcUrl(options.rpcUrlOverride),
       abiSignatures: options.abiSignatures,
-      chainConfig: this.getChainConfig(),
+      chainConfig: this.getChainConfig(options.rpcUrlOverride),
       httpProtocol: options.httpProtocol || 'https://',
       networkSpecificConfigs: options.networkSpecificConfigs,
       endpoints: this.getEndpoints(),
@@ -51,8 +52,8 @@ export abstract class BaseNetworkEnvironment<T, M> {
     return this.config.services;
   }
 
-  protected abstract getRpcUrl(): string;
-  protected abstract getChainConfig(): Chain;
+  protected abstract getRpcUrl(overrideRpc?: string): string;
+  protected abstract getChainConfig(overrideRpc?: string): Chain;
   protected abstract getEndpoints(): NagaEndpointsType;
   protected abstract getDefaultRealmId(): bigint;
 }
diff --git a/packages/networks/src/networks/vNaga/envs/naga-dev/naga-dev.env.ts b/packages/networks/src/networks/vNaga/envs/naga-dev/naga-dev.env.ts
@@ -20,7 +20,7 @@ export class NagaDevEnvironment extends BaseNetworkEnvironment<
   NagaDevSignatures,
   NagaDevSpecificConfigs
 > {
-  constructor() {
+  constructor(options?: { rpcUrlOverride?: string }) {
     super({
       network: NETWORK,
       abiSignatures: nagaDevSignatures,
@@ -34,15 +34,16 @@ export class NagaDevEnvironment extends BaseNetworkEnvironment<
       minimumThreshold: MINIMUM_THRESHOLD,
       httpProtocol: PROTOCOL,
       requiredAttestation: false,
+      rpcUrlOverride: options?.rpcUrlOverride,
     });
   }
 
-  protected getRpcUrl(): string {
-    return chainInfo.RPC_URL;
+  protected getRpcUrl(overrideRpc?: string): string {
+    return chainInfo.resolveRpcUrl(overrideRpc);
   }
 
-  protected getChainConfig(): Chain {
-    return chainInfo.viemChainConfig;
+  protected getChainConfig(overrideRpc?: string): Chain {
+    return chainInfo.buildViemChainConfig(overrideRpc);
   }
 
   protected getEndpoints(): NagaEndpointsType {
diff --git a/packages/networks/src/networks/vNaga/envs/naga-local/naga-local.env.ts b/packages/networks/src/networks/vNaga/envs/naga-local/naga-local.env.ts
@@ -21,7 +21,7 @@ export class NagaLocalEnvironment extends BaseNetworkEnvironment<
   NagaLocalSignatures,
   NagaLocalSpecificConfigs
 > {
-  constructor() {
+  constructor(options?: { rpcUrlOverride?: string }) {
     super({
       network: NETWORK,
       abiSignatures: signatures, // Note: Uses locally generated signatures
@@ -36,15 +36,25 @@ export class NagaLocalEnvironment extends BaseNetworkEnvironment<
       minimumThreshold: MINIMUM_THRESHOLD,
       httpProtocol: PROTOCOL, // Note: HTTP not HTTPS
       requiredAttestation: false,
+      rpcUrlOverride: options?.rpcUrlOverride,
     });
   }
 
-  protected getRpcUrl(): string {
-    return chainInfo.RPC_URL; // Note: Uses Anvil instead of ChronicleYellowstone
+  protected getRpcUrl(overrideRpc?: string): string {
+    return overrideRpc ?? chainInfo.RPC_URL; // Note: Uses Anvil instead of ChronicleYellowstone
   }
 
-  protected getChainConfig(): Chain {
-    return chainInfo.viemChainConfig; // Note: Anvil chain config
+  protected getChainConfig(overrideRpc?: string): Chain {
+    const rpc = overrideRpc ?? chainInfo.RPC_URL;
+    const base = chainInfo.viemChainConfig; // Note: Anvil chain config
+    return {
+      ...base,
+      rpcUrls: {
+        ...base.rpcUrls,
+        default: { ...base.rpcUrls.default, http: [rpc] },
+        public: { ...(base.rpcUrls as any)['public'], http: [rpc] },
+      },
+    } as Chain;
   }
 
   protected getEndpoints(): NagaEndpointsType {
diff --git a/packages/networks/src/networks/vNaga/envs/naga-staging/naga-staging.env.ts b/packages/networks/src/networks/vNaga/envs/naga-staging/naga-staging.env.ts
@@ -20,7 +20,7 @@ export class NagaStagingEnvironment extends BaseNetworkEnvironment<
   NagaStagingSignatures,
   NagaStagingSpecificConfigs
 > {
-  constructor() {
+  constructor(options?: { rpcUrlOverride?: string }) {
     super({
       network: NETWORK,
       abiSignatures: nagaStagingSignatures,
@@ -34,15 +34,16 @@ export class NagaStagingEnvironment extends BaseNetworkEnvironment<
       minimumThreshold: MINIMUM_THRESHOLD,
       httpProtocol: PROTOCOL,
       requiredAttestation: true,
+      rpcUrlOverride: options?.rpcUrlOverride,
     });
   }
 
-  protected getRpcUrl(): string {
-    return chainInfo.RPC_URL;
+  protected getRpcUrl(overrideRpc?: string): string {
+    return chainInfo.resolveRpcUrl(overrideRpc);
   }
 
-  protected getChainConfig(): Chain {
-    return chainInfo.viemChainConfig;
+  protected getChainConfig(overrideRpc?: string): Chain {
+    return chainInfo.buildViemChainConfig(overrideRpc);
   }
 
   protected getEndpoints(): NagaEndpointsType {
diff --git a/packages/networks/src/networks/vNaga/envs/naga-test/naga-test.env.ts b/packages/networks/src/networks/vNaga/envs/naga-test/naga-test.env.ts
@@ -20,7 +20,7 @@ export class NagaTestEnvironment extends BaseNetworkEnvironment<
   NagaTestSignatures,
   NagaTestSpecificConfigs
 > {
-  constructor() {
+  constructor(options?: { rpcUrlOverride?: string }) {
     super({
       network: NETWORK,
       abiSignatures: nagaTestSignatures,
@@ -34,15 +34,16 @@ export class NagaTestEnvironment extends BaseNetworkEnvironment<
       minimumThreshold: MINIMUM_THRESHOLD,
       httpProtocol: PROTOCOL,
       requiredAttestation: true,
+      rpcUrlOverride: options?.rpcUrlOverride,
     });
   }
 
-  protected getRpcUrl(): string {
-    return chainInfo.RPC_URL;
+  protected getRpcUrl(overrideRpc?: string): string {
+    return chainInfo.resolveRpcUrl(overrideRpc);
   }
 
-  protected getChainConfig(): Chain {
-    return chainInfo.viemChainConfig;
+  protected getChainConfig(overrideRpc?: string): Chain {
+    return chainInfo.buildViemChainConfig(overrideRpc);
   }
 
   protected getEndpoints(): NagaEndpointsType {
diff --git a/packages/networks/src/networks/vNaga/shared/factories/BaseModuleFactory.ts b/packages/networks/src/networks/vNaga/shared/factories/BaseModuleFactory.ts
@@ -16,6 +16,7 @@ import { z } from 'zod';
 import { LitNetworkModuleBase } from '../../../types';
 import type { ExpectedAccountOrWalletClient } from '../managers/contract-manager/createContractsManager';
 import type { INetworkConfig } from '../interfaces/NetworkContext';
+import { createChainManagerFactory } from './BaseChainManagerFactory';
 
 // Shared utilities
 import { NetworkError } from '@lit-protocol/constants';
@@ -1102,56 +1103,38 @@ export function createBaseModule<T, M>(config: BaseModuleConfig<T, M>) {
     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,
-          });
+      // Build an overridden network config and a chain manager bound to it
+      const overriddenChainConfig = {
+        ...networkConfig.chainConfig,
+        rpcUrls: {
+          ...networkConfig.chainConfig.rpcUrls,
+          default: {
+            ...networkConfig.chainConfig.rpcUrls.default,
+            http: [resolvedRpcUrl],
+          },
+          ['public']: {
+            ...(networkConfig.chainConfig.rpcUrls as any)['public'],
+            http: [resolvedRpcUrl],
+          },
         },
-      } as typeof baseModule & {
-        withOverrides: (o: { rpcUrl?: string }) => any;
-      };
-
-      return wrapped;
+      } as typeof networkConfig.chainConfig;
+
+      const overriddenNetworkConfig = {
+        ...networkConfig,
+        rpcUrl: resolvedRpcUrl,
+        chainConfig: overriddenChainConfig,
+      } as typeof networkConfig;
+
+      const createChainManagerOverridden = (account: ExpectedAccountOrWalletClient) =>
+        createChainManagerFactory(overriddenNetworkConfig, account);
+
+      // Rebuild a fresh module bound to the overridden config
+      return createBaseModule({
+        networkConfig: overriddenNetworkConfig,
+        moduleName,
+        createChainManager: createChainManagerOverridden,
+        verifyReleaseId: baseModule.getVerifyReleaseId(),
+      });
     },
   };
 
