feat(wrapped-keys): LIT-3956 - Use new service method storePrivateKeyBatch() to write keys to store backend in batchGeneratePrivateKeys() client API

Author: MaximusHaximus

packages/wrapped-keys/src/lib/api/batch-generate-private-keys.ts

@@ -1,7 +1,7 @@
 import { getKeyTypeFromNetwork } from './utils';
 import { batchGenerateKeysWithLitAction } from '../lit-actions-client';
 import { getLitActionCodeOrCidCommon } from '../lit-actions-client/utils';
-import { storePrivateKey } from '../service-client';
+import { storePrivateKeyBatch } from '../service-client';
 import {
   BatchGeneratePrivateKeysActionResult,
   BatchGeneratePrivateKeysParams,
@@ -41,34 +41,40 @@ export async function batchGeneratePrivateKeys(
     pkpSessionSigs,
   });
 
-  const results = await Promise.all(
-    actionResults.map(
-      async (result): Promise<BatchGeneratePrivateKeysActionResult> => {
-        const { generateEncryptedPrivateKey, network } = result;
+  const keyParamsBatch = actionResults.map((keyData) => {
+    const { generateEncryptedPrivateKey, network } = keyData;
+    return {
+      ...generateEncryptedPrivateKey,
+      keyType: getKeyTypeFromNetwork(network),
+    };
+  });
+
+  const { ids } = await storePrivateKeyBatch({
+    pkpAddress,
+    sessionSig,
+    storedKeyMetadataBatch: keyParamsBatch,
+    litNetwork: litNodeClient.config.litNetwork,
+  });
 
-        const signature = result.signMessage?.signature;
+  const results = actionResults.map(
+    (actionResult, ndx): BatchGeneratePrivateKeysActionResult => {
+      const {
+        generateEncryptedPrivateKey: { memo, publicKey },
+      } = actionResult;
+      const id = ids[ndx]; // Result of writes is in same order as provided
 
-        const { id } = await storePrivateKey({
-          sessionSig,
-          storedKeyMetadata: {
-            ...generateEncryptedPrivateKey,
-            keyType: getKeyTypeFromNetwork(network),
-            pkpAddress,
-          },
-          litNetwork: litNodeClient.config.litNetwork,
-        });
+      const signature = actionResult.signMessage?.signature;
 
-        return {
-          ...(signature ? { signMessage: { signature } } : {}),
-          generateEncryptedPrivateKey: {
-            memo: generateEncryptedPrivateKey.memo,
-            id,
-            generatedPublicKey: generateEncryptedPrivateKey.publicKey,
-            pkpAddress,
-          },
-        };
-      }
-    )
+      return {
+        ...(signature ? { signMessage: { signature } } : {}),
+        generateEncryptedPrivateKey: {
+          memo: memo,
+          id,
+          generatedPublicKey: publicKey,
+          pkpAddress,
+        },
+      };
+    }
   );
 
   return { pkpAddress, results };

packages/wrapped-keys/src/lib/api/store-encrypted-key-batch.ts

@@ -46,6 +46,7 @@ export async function storeEncryptedKeyBatch(
     );
 
   return storePrivateKeyBatch({
+    pkpAddress,
     storedKeyMetadataBatch,
     sessionSig: getFirstSessionSig(pkpSessionSigs),
     litNetwork: litNodeClient.config.litNetwork,

packages/wrapped-keys/src/lib/service-client/types.ts

@@ -31,14 +31,10 @@ export interface StoreKeyParams extends BaseApiParams {
 }
 
 export interface StoreKeyBatchParams extends BaseApiParams {
+  pkpAddress: string;
   storedKeyMetadataBatch: Pick<
     StoredKeyData,
-    | 'pkpAddress'
-    | 'publicKey'
-    | 'keyType'
-    | 'dataToEncryptHash'
-    | 'ciphertext'
-    | 'memo'
+    'publicKey' | 'keyType' | 'dataToEncryptHash' | 'ciphertext' | 'memo'
   >[];
 }