Merge pull request #839 from LIT-Protocol/feature/node-4604-fix-litclient-return-type

Feature/node 4604 fix litclient return type

Author: Ansonhkg

e2e/artillery/processors/multi-endpoint.ts

@@ -38,29 +38,42 @@ import {
 let sharedContext: any = null;
 let alicePkpAuthContext: any = null;
 let aliceCustomAuthContext: any = null;
+let initializationPromise: Promise<any> | null = null;
 
 /**
  * Initialize the shared context once per Artillery run
  */
 async function initializeSharedContext() {
   if (sharedContext) return sharedContext;
+  
+  // Prevent race conditions by ensuring only one initialization happens
+  if (initializationPromise) {
+    return await initializationPromise;
+  }
 
-  try {
-    console.log('🚀 Initializing Artillery shared context...');
-
-    // Use the same init function as e2e tests
-    sharedContext = await init();
-
-    // Create auth contexts using helper functions
-    alicePkpAuthContext = await createPkpAuthContext(sharedContext);
-    aliceCustomAuthContext = await createCustomAuthContext(sharedContext);
+  initializationPromise = (async () => {
+    try {
+      console.log('🚀 Initializing Artillery shared context...');
+
+      // Use the same init function as e2e tests
+      sharedContext = await init();
+
+      // Create auth contexts using helper functions
+      alicePkpAuthContext = await createPkpAuthContext(sharedContext);
+      aliceCustomAuthContext = await createCustomAuthContext(sharedContext);
+
+      console.log('✅ Artillery shared context initialized');
+      return sharedContext;
+    } catch (error) {
+      console.error('❌ Failed to initialize Artillery context:', error);
+      // Reset state on failure so retry is possible
+      initializationPromise = null;
+      sharedContext = null;
+      throw error;
+    }
+  })();
 
-    console.log('✅ Artillery shared context initialized');
-    return sharedContext;
-  } catch (error) {
-    console.error('❌ Failed to initialize Artillery context:', error);
-    throw error;
-  }
+  return await initializationPromise;
 }
 
 /**
@@ -166,7 +179,7 @@ export async function runMultiEndpointTest(context: any, events: any) {
 /**
  * PKP Sign test functions
  */
-export async function runPkpSignTest(context, events) {
+export async function runPkpSignTest(context: any, events: any) {
   await initializeSharedContext();
 
   const parallelism = context.vars.parallelism || 5;
@@ -176,7 +189,7 @@ export async function runPkpSignTest(context, events) {
   await runTestWithMetrics('pkp_sign', testFn, context, events, parallelism);
 }
 
-export async function runPkpSignTestWithEoa(context, events) {
+export async function runPkpSignTestWithEoa(context: any, events: any) {
   await initializeSharedContext();
 
   const parallelism = context.vars.parallelism || 5;
@@ -192,7 +205,7 @@ export async function runPkpSignTestWithEoa(context, events) {
   );
 }
 
-export async function runPkpSignTestWithPkp(context, events) {
+export async function runPkpSignTestWithPkp(context: any, events: any) {
   await initializeSharedContext();
 
   const parallelism = context.vars.parallelism || 5;
@@ -208,7 +221,7 @@ export async function runPkpSignTestWithPkp(context, events) {
   );
 }
 
-export async function runPkpSignTestWithCustom(context, events) {
+export async function runPkpSignTestWithCustom(context: any, events: any) {
   await initializeSharedContext();
 
   const parallelism = context.vars.parallelism || 5;
@@ -227,7 +240,7 @@ export async function runPkpSignTestWithCustom(context, events) {
 /**
  * Encrypt/Decrypt test functions
  */
-export async function runEncryptDecryptTest(context, events) {
+export async function runEncryptDecryptTest(context: any, events: any) {
   await initializeSharedContext();
 
   const parallelism = context.vars.parallelism || 3;
@@ -243,7 +256,7 @@ export async function runEncryptDecryptTest(context, events) {
   );
 }
 
-export async function runPkpEncryptDecryptTest(context, events) {
+export async function runPkpEncryptDecryptTest(context: any, events: any) {
   await initializeSharedContext();
 
   const parallelism = context.vars.parallelism || 3;
@@ -259,7 +272,7 @@ export async function runPkpEncryptDecryptTest(context, events) {
   );
 }
 
-export async function runEncryptDecryptFlowTest(context, events) {
+export async function runEncryptDecryptFlowTest(context: any, events: any) {
   await initializeSharedContext();
 
   const parallelism = context.vars.parallelism || 3;
@@ -278,7 +291,7 @@ export async function runEncryptDecryptFlowTest(context, events) {
 /**
  * Execute JS test function
  */
-export async function runExecuteJsTest(context, events) {
+export async function runExecuteJsTest(context: any, events: any) {
   await initializeSharedContext();
 
   const parallelism = context.vars.parallelism || 4;
@@ -291,7 +304,7 @@ export async function runExecuteJsTest(context, events) {
 /**
  * View PKPs test functions
  */
-export async function runViewPkpsTest(context, events) {
+export async function runViewPkpsTest(context: any, events: any) {
   await initializeSharedContext();
 
   const parallelism = context.vars.parallelism || 5;
@@ -325,7 +338,7 @@ export async function runViewPkpsTest(context, events) {
 /**
  * Viem integration test functions
  */
-export async function runViemSignTest(context, events) {
+export async function runViemSignTest(context: any, events: any) {
   await initializeSharedContext();
 
   const parallelism = context.vars.parallelism || 3;

e2e/src/init.ts

@@ -50,7 +50,7 @@ export const init = async (
   const aliceViemAccountAuthData = await ViemAccountAuthenticator.authenticate(
     aliceViemAccount
   );
-  
+
   const bobViemAccount = privateKeyToAccount(generatePrivateKey());
   const bobViemAccountAuthData = await ViemAccountAuthenticator.authenticate(
     bobViemAccount
@@ -167,7 +167,7 @@ export const init = async (
     },
     storageProvider: storagePlugins.localStorageNode({
       appName: 'my-app',
-      networkName: 'naga-dev',
+      networkName: _network,
       storagePath: './pkp-tokens',
     }),
   });
@@ -180,7 +180,7 @@ export const init = async (
     },
     storageProvider: storagePlugins.localStorageNode({
       appName: 'my-app',
-      networkName: 'naga-dev',
+      networkName: _network,
       storagePath: './pkp-tokens-bob',
     }),
   });
@@ -219,7 +219,7 @@ export const init = async (
     },
     storageProvider: storagePlugins.localStorageNode({
       appName: 'my-app',
-      networkName: 'naga-dev',
+      networkName: _network,
       storagePath: './pkp-tokens',
     }),
   });
@@ -232,7 +232,7 @@ export const init = async (
     },
     storageProvider: storagePlugins.localStorageNode({
       appName: 'my-app',
-      networkName: 'naga-dev',
+      networkName: _network,
       storagePath: './pkp-tokens-bob',
     }),
   });

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

@@ -16,6 +16,7 @@ import type {
 } from '@lit-protocol/networks';
 import {
   AuthContextSchema2,
+  AuthData,
   EncryptedVersion1Schema,
   HexPrefixedSchema,
   JsonSignCustomSessionKeyRequestForPkpReturnSchema,
@@ -50,8 +51,7 @@ import {
   MintWithCustomAuthRequest,
   MintWithCustomAuthSchema,
 } from './schemas/MintWithCustomAuthSchema';
-import { LitClient, NagaNetworkModule } from './type';
-import { NagaLitClient } from './types/NagaLitClient.type';
+import { NagaNetworkModule } from './type';
 
 const _logger = getChildLogger({
   module: 'createLitClient',
@@ -84,7 +84,7 @@ export const createLitClient = async ({
   network,
 }: {
   network: SupportedNetworkModule;
-}): Promise<LitClient> => {
+}) => {
   switch (network.id) {
     // -- (v8) Naga Network Module
     case 'naga':
@@ -144,7 +144,7 @@ export const createLitClient = async ({
  */
 export const _createNagaLitClient = async (
   networkModule: NagaNetworkModule
-): Promise<NagaLitClient> => {
+) => {
   const _stateManager = await networkModule.createStateManager<
     Awaited<ReturnType<typeof orchestrateHandshake>>,
     NagaNetworkModule
@@ -702,7 +702,7 @@ export const _createNagaLitClient = async (
     return response;
   }
 
-  const litClient: NagaLitClient = {
+  const litClient = {
     // This function is likely be used by another module to get the current context, eg. auth manager
     // only adding what is required by other modules for now.
     // maybe you will need connectionInfo: _stateManager.getLatestConnectionInfo(),
@@ -832,7 +832,7 @@ export const _createNagaLitClient = async (
         authMethodType: number | bigint;
         authMethodId: string;
         accessToken?: string;
-      };
+      } | AuthData;
       pagination?: { limit?: number; offset?: number };
       storageProvider?: PKPStorageProvider;
     }) => {

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

@@ -2,7 +2,7 @@
 import { NagaLocalModule } from '@lit-protocol/networks';
 import { NagaDevModule } from '@lit-protocol/networks';
 import { NagaStagingModule } from '@lit-protocol/networks';
-import { NagaLitClient } from './types/NagaLitClient.type';
+// import { NagaLitClient } from './types/NagaLitClient.type';
 
 /**
  * ========== All Network Modules ==========
@@ -27,5 +27,5 @@ export type NagaNetworkModule =
 /**
  * Union type for all possible Lit clients
  */
-export type LitClient = NagaLitClient;
+// export type LitClient = NagaLitClient;
 // | DatilLitClient;