@fhevm/sdk 1.1.0-alpha.4 → 1.1.0-alpha.5

package/core/kms/kmsUserDecryptEip712V2Types.ts

@@ -0,0 +1,31 @@
+import type { KmsUserDecryptEip712V2Types } from '../types/kms.js';
+
+////////////////////////////////////////////////////////////////////////////////
+// KmsUserDecryptEIP712V2Types (RFC-016)
+////////////////////////////////////////////////////////////////////////////////
+
+export const kmsUserDecryptEip712V2Types: KmsUserDecryptEip712V2Types = {
+  EIP712Domain: [
+    { name: 'name', type: 'string' },
+    { name: 'version', type: 'string' },
+    { name: 'chainId', type: 'uint256' },
+    { name: 'verifyingContract', type: 'address' },
+  ] as const,
+  // CRITICAL: Field order is authoritative — it determines the EIP-712 type hash.
+  // Changing the order will produce a different signature and break on-chain verification.
+  // Must match the Solidity struct definition exactly.
+  UserDecryptRequestVerification: [
+    { name: 'userAddress', type: 'address' },
+    { name: 'publicKey', type: 'bytes' },
+    { name: 'allowedContracts', type: 'address[]' },
+    { name: 'startTimestamp', type: 'uint256' },
+    { name: 'durationSeconds', type: 'uint256' },
+    { name: 'extraData', type: 'bytes' },
+  ] as const,
+} as const;
+
+Object.freeze(kmsUserDecryptEip712V2Types);
+Object.freeze(kmsUserDecryptEip712V2Types.EIP712Domain);
+Object.freeze(kmsUserDecryptEip712V2Types.UserDecryptRequestVerification);
+kmsUserDecryptEip712V2Types.EIP712Domain.forEach(Object.freeze);
+kmsUserDecryptEip712V2Types.UserDecryptRequestVerification.forEach(Object.freeze);

package/core/kms/publicDecrypt.ts

@@ -39,7 +39,8 @@ export async function publicDecrypt(context: Context, parameters: Parameters): P
   // EIP-712 signature from the user, so the SDK can safely fetch the current
   // context ID and build the extraData transparently to the user.
   const requestedKmsSignersContext = await readKmsSignersContext(context, {
-    address: context.chain.fhevm.contracts.kmsVerifier.address as ChecksummedAddress,
+    kmsVerifierAddress: context.chain.fhevm.contracts.kmsVerifier.address as ChecksummedAddress,
+    protocolConfigAddress: context.chain.fhevm.contracts.protocolConfig?.address as ChecksummedAddress | undefined,
   });
 
   const requestedExtraData: BytesHex = kmsSignersContextToExtraData(requestedKmsSignersContext);
@@ -65,7 +66,7 @@ export async function publicDecrypt(context: Context, parameters: Parameters): P
 
   // 4. Check: ACL permissions
   await checkAllowedForDecryption(context, {
-    address: context.chain.fhevm.contracts.acl.address as ChecksummedAddress,
+    aclAddress: context.chain.fhevm.contracts.acl.address as ChecksummedAddress,
     handles: orderedHandles,
   });
 
@@ -84,7 +85,8 @@ export async function publicDecrypt(context: Context, parameters: Parameters): P
 
   // 6. Reconcile KMS signer context using 'loose' mode
   const reconciledKmsSignersContext: KmsSignersContext = await reconcileKmsSignersContext(context, {
-    address: context.chain.fhevm.contracts.kmsVerifier.address as ChecksummedAddress,
+    kmsVerifierAddress: context.chain.fhevm.contracts.kmsVerifier.address as ChecksummedAddress,
+    protocolConfigAddress: context.chain.fhevm.contracts.protocolConfig?.address as ChecksummedAddress | undefined,
     relayerExtraData,
     requestedKmsSignersContext: requestedKmsSignersContext,
     mode: 'loose',

package/core/modules/decrypt/mock.ts

@@ -2,7 +2,6 @@ import type { FhevmRuntime } from '../../types/coreFhevmRuntime.js';
 import type {
   DecryptAndReconstructParameters,
   DecryptAndReconstructReturnType,
-  DecryptAndReconstructUserParameters,
   DeserializeTkmsPrivateKeyParameters,
   DeserializeTkmsPrivateKeyReturnType,
   GenerateTkmsPrivateKeyReturnType,
@@ -11,9 +10,7 @@ import type {
   SerializeTkmsPrivateKeyParameters,
   SerializeTkmsPrivateKeyReturnType,
   VerifyTkmsPrivateKeyParameters,
-  UserDecryptModuleParameters,
   DecryptModuleFactory,
-  UserDecryptModuleFactory,
 } from './types.js';
 import type { CleartextEthereumModule } from '../ethereum/types-ct.js';
 import type { TkmsPrivateKey } from '../../types/tkms-p.js';
@@ -213,35 +210,3 @@ export const decryptModule: DecryptModuleFactory = (runtime: FhevmRuntime) => {
     }),
   });
 };
-
-//////////////////////////////////////////////////////////////////////////////
-// userDecryptActions
-//////////////////////////////////////////////////////////////////////////////
-
-export const userDecryptModule: UserDecryptModuleFactory = (
-  runtime: FhevmRuntime,
-  parameters: UserDecryptModuleParameters,
-) => {
-  const { privateKey } = parameters;
-  return Object.freeze({
-    userDecrypt: Object.freeze({
-      initTkmsModule: () => Promise.resolve(),
-      getTkmsModuleInfo: () => {
-        throw new Error('Not yet implemented');
-      },
-      decryptAndReconstruct: (args: DecryptAndReconstructUserParameters) =>
-        decryptAndReconstruct(runtime, {
-          ...args,
-          tkmsPrivateKey: privateKey,
-        }),
-      getTkmsPublicKeyHex: () =>
-        getTkmsPublicKeyHex(runtime, {
-          tkmsPrivateKey: privateKey,
-        }),
-      serializeTkmsPrivateKey: () =>
-        serializeTkmsPrivateKey(runtime, {
-          tkmsPrivateKey: privateKey,
-        }),
-    }),
-  });
-};

package/core/modules/decrypt/module/api-p.ts

@@ -5,6 +5,7 @@ import type {
   Client,
   TypedPlaintext,
   KmsLibApi,
+  TkmsVersion,
 } from '../../../../wasm/tkms/KmsLibApi.js';
 import type { FhevmRuntime } from '../../../types/coreFhevmRuntime.js';
 import type { ClearValue } from '../../../types/encryptedTypes-p.js';
@@ -15,6 +16,7 @@ import type {
   DecryptAndReconstructReturnType,
   DeserializeTkmsPrivateKeyParameters,
   DeserializeTkmsPrivateKeyReturnType,
+  GenerateTkmsPrivateKeyParameters,
   GenerateTkmsPrivateKeyReturnType,
   GetTkmsPublicKeyHexParameters,
   GetTkmsPublicKeyHexReturnType,
@@ -95,29 +97,49 @@ function verifyToken(token: symbol): void {
 class TkmsPublicEncKeyMlKem512Impl {
   readonly #publicEncKeyMlKem512Wasm: PublicEncKeyMlKem512WasmType;
   #bytesHex: BytesHex | undefined;
+  readonly #tkmsVersion: TkmsVersion;
 
-  constructor(token: symbol, publicEncKeyMlKem512Wasm: PublicEncKeyMlKem512WasmType) {
+  constructor(token: symbol, tkmsVersion: TkmsVersion, publicEncKeyMlKem512Wasm: PublicEncKeyMlKem512WasmType) {
     verifyToken(token);
     this.#publicEncKeyMlKem512Wasm = publicEncKeyMlKem512Wasm;
+    this.#tkmsVersion = tkmsVersion;
   }
 
-  public static [GET_NATIVE_FUNC](key: unknown, token: symbol): PublicEncKeyMlKem512WasmType {
+  public get tkmsVersion(): TkmsVersion {
+    return this.#tkmsVersion;
+  }
+
+  public static [GET_NATIVE_FUNC](key: unknown, token: symbol, tkmsVersion: TkmsVersion): PublicEncKeyMlKem512WasmType {
     verifyToken(token);
     if (!(key instanceof TkmsPublicEncKeyMlKem512Impl)) {
       throw new Error('Unauthorized');
     }
+    if (key.#tkmsVersion !== tkmsVersion) {
+      throw new Error(`TkmsVersion mismatch`);
+    }
     return key.#publicEncKeyMlKem512Wasm;
   }
 
-  public static [GET_BYTES_HEX_FUNC](key: unknown, token: symbol, kmsLibApi: KmsLibApi): BytesHex {
+  public static [GET_BYTES_HEX_FUNC](
+    key: unknown,
+    token: symbol,
+    tkmsVersion: TkmsVersion,
+    kmsLibApi: KmsLibApi,
+  ): BytesHex {
     verifyToken(token);
     if (!(key instanceof TkmsPublicEncKeyMlKem512Impl)) {
       throw new Error('Unauthorized');
     }
+
+    if (key.#tkmsVersion !== tkmsVersion || kmsLibApi.getWasmInfo().version !== tkmsVersion) {
+      throw new Error(`TkmsVersion mismatch`);
+    }
+
     if (key.#bytesHex === undefined) {
       const bytes: Bytes = kmsLibApi.ml_kem_pke_pk_to_u8vec(key.#publicEncKeyMlKem512Wasm);
       key.#bytesHex = bytesToHexLarge(bytes, false /* no 0x */);
     }
+
     return key.#bytesHex;
   }
 }
@@ -131,31 +153,55 @@ class TkmsPrivateEncKeyMlKem512Impl implements TkmsPrivateKey {
 
   readonly #privateEncKeyMlKem512Wasm: PrivateEncKeyMlKem512WasmType;
   #publicKey: TkmsPublicEncKeyMlKem512Impl | undefined;
+  readonly #tkmsVersion: TkmsVersion;
 
-  constructor(token: symbol, privateEncKeyMlKem512Wasm: PrivateEncKeyMlKem512WasmType) {
+  constructor(token: symbol, tkmsVersion: TkmsVersion, privateEncKeyMlKem512Wasm: PrivateEncKeyMlKem512WasmType) {
     verifyToken(token);
     this.#privateEncKeyMlKem512Wasm = privateEncKeyMlKem512Wasm;
+    this.#tkmsVersion = tkmsVersion;
   }
 
-  public static [GET_NATIVE_FUNC](key: unknown, token: symbol): PrivateEncKeyMlKem512WasmType {
+  public get tkmsVersion(): TkmsVersion {
+    return this.#tkmsVersion;
+  }
+
+  public static [GET_NATIVE_FUNC](
+    key: unknown,
+    token: symbol,
+    tkmsVersion: TkmsVersion,
+  ): PrivateEncKeyMlKem512WasmType {
     verifyToken(token);
     if (!(key instanceof TkmsPrivateEncKeyMlKem512Impl)) {
       throw new Error('Unauthorized');
     }
+    if (key.#tkmsVersion !== tkmsVersion) {
+      throw new Error(`TkmsVersion mismatch`);
+    }
     return key.#privateEncKeyMlKem512Wasm;
   }
 
-  public static [GET_PUBLIC_KEY_FUNC](key: unknown, token: symbol, kmsLibApi: KmsLibApi): TkmsPublicEncKeyMlKem512Impl {
+  public static [GET_PUBLIC_KEY_FUNC](
+    key: unknown,
+    token: symbol,
+    tkmsVersion: TkmsVersion,
+    kmsLibApi: KmsLibApi,
+  ): TkmsPublicEncKeyMlKem512Impl {
     verifyToken(token);
     if (!(key instanceof TkmsPrivateEncKeyMlKem512Impl)) {
       throw new Error('Unauthorized');
     }
+
+    if (key.#tkmsVersion !== tkmsVersion || kmsLibApi.getWasmInfo().version !== tkmsVersion) {
+      throw new Error(`TkmsVersion mismatch`);
+    }
+
     if (key.#publicKey === undefined) {
       const publicEncKeyMlKem512Wasm = kmsLibApi.ml_kem_pke_get_pk(
         key.#privateEncKeyMlKem512Wasm,
       ) as PublicEncKeyMlKem512WasmType;
-      key.#publicKey = new TkmsPublicEncKeyMlKem512Impl(token, publicEncKeyMlKem512Wasm);
+      key.#publicKey = new TkmsPublicEncKeyMlKem512Impl(token, key.#tkmsVersion, publicEncKeyMlKem512Wasm);
     }
+
     return key.#publicKey;
   }
 }
@@ -164,13 +210,16 @@ class TkmsPrivateEncKeyMlKem512Impl implements TkmsPrivateKey {
 // generateTkmsPrivateKey
 //////////////////////////////////////////////////////////////////////////////
 
-export async function generateTkmsPrivateKey(runtime: FhevmRuntime): Promise<GenerateTkmsPrivateKeyReturnType> {
-  const kmsLib = await initTkmsModule(runtime);
+export async function generateTkmsPrivateKey(
+  runtime: FhevmRuntime,
+  parameters: GenerateTkmsPrivateKeyParameters,
+): Promise<GenerateTkmsPrivateKeyReturnType> {
+  const kmsLib = await initTkmsModule(runtime, { tkmsVersion: parameters.tkmsVersion });
 
   const privateEncKeyMlKem512Wasm: PrivateEncKeyMlKem512WasmType =
     kmsLib.ml_kem_pke_keygen() as PrivateEncKeyMlKem512WasmType;
 
-  return new TkmsPrivateEncKeyMlKem512Impl(PRIVATE_TKMS_LIB_TOKEN, privateEncKeyMlKem512Wasm);
+  return new TkmsPrivateEncKeyMlKem512Impl(PRIVATE_TKMS_LIB_TOKEN, parameters.tkmsVersion, privateEncKeyMlKem512Wasm);
 }
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -181,7 +230,7 @@ export async function decryptAndReconstruct(
   runtime: FhevmRuntime,
   parameters: DecryptAndReconstructParameters,
 ): Promise<DecryptAndReconstructReturnType> {
-  const kmsLib = await initTkmsModule(runtime);
+  const kmsLib = await initTkmsModule(runtime, { tkmsVersion: parameters.tkmsVersion });
 
   const { tkmsPrivateKey, shares } = parameters;
   if (!(tkmsPrivateKey instanceof TkmsPrivateEncKeyMlKem512Impl)) {
@@ -191,10 +240,15 @@ export async function decryptAndReconstruct(
   const tkmsPublicKey: TkmsPublicEncKeyMlKem512Impl = TkmsPrivateEncKeyMlKem512Impl[GET_PUBLIC_KEY_FUNC](
     tkmsPrivateKey,
     PRIVATE_TKMS_LIB_TOKEN,
+    parameters.tkmsVersion,
     kmsLib,
   );
 
   const metadata: KmsSigncryptedSharesMetadata = getMetadata(shares);
+  if (metadata.tkmsVersion !== parameters.tkmsVersion) {
+    throw new Error('TkmsVersion mismatch');
+  }
+
   const sharesArray: readonly KmsSigncryptedShare[] = getShares(shares);
 
   const firstShare = sharesArray[0];
@@ -227,16 +281,19 @@ export async function decryptAndReconstruct(
   const privateEncKeyMlKem512Wasm: PrivateEncKeyMlKem512WasmType = TkmsPrivateEncKeyMlKem512Impl[GET_NATIVE_FUNC](
     tkmsPrivateKey,
     PRIVATE_TKMS_LIB_TOKEN,
+    parameters.tkmsVersion,
   );
 
   const publicEncKeyMlKem512Wasm: PublicEncKeyMlKem512WasmType = TkmsPublicEncKeyMlKem512Impl[GET_NATIVE_FUNC](
     tkmsPublicKey,
     PRIVATE_TKMS_LIB_TOKEN,
+    parameters.tkmsVersion,
   );
 
   const publicEncKeyMlKem512WasmBytesHex: BytesHex = TkmsPublicEncKeyMlKem512Impl[GET_BYTES_HEX_FUNC](
     tkmsPublicKey,
     PRIVATE_TKMS_LIB_TOKEN,
+    parameters.tkmsVersion,
     kmsLib,
   );
 
@@ -339,7 +396,7 @@ export async function getTkmsPublicKeyHex(
   runtime: FhevmRuntime,
   parameters: GetTkmsPublicKeyHexParameters,
 ): Promise<GetTkmsPublicKeyHexReturnType> {
-  const kmsLib = await initTkmsModule(runtime);
+  const kmsLib = await initTkmsModule(runtime, { tkmsVersion: parameters.tkmsVersion });
 
   const { tkmsPrivateKey } = parameters;
 
@@ -347,9 +404,19 @@ export async function getTkmsPublicKeyHex(
     throw new Error('Invalid tkmsPrivateKey');
   }
 
-  const publicKey = TkmsPrivateEncKeyMlKem512Impl[GET_PUBLIC_KEY_FUNC](tkmsPrivateKey, PRIVATE_TKMS_LIB_TOKEN, kmsLib);
+  const publicKey = TkmsPrivateEncKeyMlKem512Impl[GET_PUBLIC_KEY_FUNC](
+    tkmsPrivateKey,
+    PRIVATE_TKMS_LIB_TOKEN,
+    parameters.tkmsVersion,
+    kmsLib,
+  );
 
-  return TkmsPublicEncKeyMlKem512Impl[GET_BYTES_HEX_FUNC](publicKey, PRIVATE_TKMS_LIB_TOKEN, kmsLib);
+  return TkmsPublicEncKeyMlKem512Impl[GET_BYTES_HEX_FUNC](
+    publicKey,
+    PRIVATE_TKMS_LIB_TOKEN,
+    parameters.tkmsVersion,
+    kmsLib,
+  );
 }
 
 //////////////////////////////////////////////////////////////////////////////
@@ -360,7 +427,7 @@ export async function serializeTkmsPrivateKey(
   runtime: FhevmRuntime,
   parameters: SerializeTkmsPrivateKeyParameters,
 ): Promise<SerializeTkmsPrivateKeyReturnType> {
-  const kmsLib = await initTkmsModule(runtime);
+  const kmsLib = await initTkmsModule(runtime, { tkmsVersion: parameters.tkmsVersion });
 
   const { tkmsPrivateKey } = parameters;
 
@@ -371,6 +438,7 @@ export async function serializeTkmsPrivateKey(
   const privateEncKeyMlKem512Wasm: PrivateEncKeyMlKem512WasmType = TkmsPrivateEncKeyMlKem512Impl[GET_NATIVE_FUNC](
     tkmsPrivateKey,
     PRIVATE_TKMS_LIB_TOKEN,
+    parameters.tkmsVersion,
   );
 
   return kmsLib.ml_kem_pke_sk_to_u8vec(privateEncKeyMlKem512Wasm);
@@ -384,7 +452,7 @@ export async function deserializeTkmsPrivateKey(
   runtime: FhevmRuntime,
   parameters: DeserializeTkmsPrivateKeyParameters,
 ): Promise<DeserializeTkmsPrivateKeyReturnType> {
-  const kmsLib = await initTkmsModule(runtime);
+  const kmsLib = await initTkmsModule(runtime, { tkmsVersion: parameters.tkmsVersion });
 
   const { tkmsPrivateKeyBytes } = parameters;
 
@@ -392,7 +460,7 @@ export async function deserializeTkmsPrivateKey(
     tkmsPrivateKeyBytes,
   ) as PrivateEncKeyMlKem512WasmType;
 
-  return new TkmsPrivateEncKeyMlKem512Impl(PRIVATE_TKMS_LIB_TOKEN, privateEncKeyMlKem512Wasm);
+  return new TkmsPrivateEncKeyMlKem512Impl(PRIVATE_TKMS_LIB_TOKEN, parameters.tkmsVersion, privateEncKeyMlKem512Wasm);
 }
 
 //////////////////////////////////////////////////////////////////////////////
@@ -403,4 +471,7 @@ export function verifyTkmsPrivateKey(_runtime: FhevmRuntime, parameters: VerifyT
   if (!(parameters.tkmsPrivateKey instanceof TkmsPrivateEncKeyMlKem512Impl)) {
     throw new Error('Invalid TkmsPrivateKey');
   }
+  if (parameters.tkmsVersion !== parameters.tkmsPrivateKey.tkmsVersion) {
+    throw new Error(`TkmsVersion mismatch`);
+  }
 }

package/core/modules/decrypt/module/index.ts

@@ -1,13 +1,16 @@
 import type { FhevmRuntime } from '../../../types/coreFhevmRuntime.js';
 import type {
   DecryptAndReconstructParameters,
-  DecryptAndReconstructUserParameters,
+  //DecryptAndReconstructUserParameters,
   DecryptModuleFactory,
   DeserializeTkmsPrivateKeyParameters,
+  GenerateTkmsPrivateKeyParameters,
+  GetTkmsModuleInfoParameters,
   GetTkmsPublicKeyHexParameters,
+  InitTkmsModuleParameters,
   SerializeTkmsPrivateKeyParameters,
-  UserDecryptModuleFactory,
-  UserDecryptModuleParameters,
+  //UserDecryptModuleFactory,
+  //UserDecryptModuleParameters,
   VerifyTkmsPrivateKeyParameters,
 } from '../types.js';
 import {
@@ -27,12 +30,12 @@ import { getTkmsModuleInfo, initTkmsModule } from './init-p.js';
 export const decryptModule: DecryptModuleFactory = (runtime: FhevmRuntime) => {
   return Object.freeze({
     decrypt: Object.freeze({
-      initTkmsModule: async () => {
-        await initTkmsModule(runtime);
+      initTkmsModule: async (args: InitTkmsModuleParameters) => {
+        await initTkmsModule(runtime, args);
       },
-      getTkmsModuleInfo: () => getTkmsModuleInfo(),
+      getTkmsModuleInfo: (args: GetTkmsModuleInfoParameters) => getTkmsModuleInfo(args),
       decryptAndReconstruct: (args: DecryptAndReconstructParameters) => decryptAndReconstruct(runtime, args),
-      generateTkmsPrivateKey: () => generateTkmsPrivateKey(runtime),
+      generateTkmsPrivateKey: (args: GenerateTkmsPrivateKeyParameters) => generateTkmsPrivateKey(runtime, args),
       serializeTkmsPrivateKey: (args: SerializeTkmsPrivateKeyParameters) => serializeTkmsPrivateKey(runtime, args),
       deserializeTkmsPrivateKey: (args: DeserializeTkmsPrivateKeyParameters) =>
         deserializeTkmsPrivateKey(runtime, args),
@@ -48,30 +51,30 @@ export const decryptModule: DecryptModuleFactory = (runtime: FhevmRuntime) => {
 // userDecryptModule
 //////////////////////////////////////////////////////////////////////////////
 
-export const userDecryptModule: UserDecryptModuleFactory = (
-  runtime: FhevmRuntime,
-  parameters: UserDecryptModuleParameters,
-) => {
-  const { privateKey } = parameters;
-  return Object.freeze({
-    userDecrypt: Object.freeze({
-      initTkmsModule: async () => {
-        await initTkmsModule(runtime);
-      },
-      getTkmsModuleInfo: () => getTkmsModuleInfo(),
-      decryptAndReconstruct: (args: DecryptAndReconstructUserParameters) =>
-        decryptAndReconstruct(runtime, {
-          ...args,
-          tkmsPrivateKey: privateKey,
-        }),
-      getTkmsPublicKeyHex: () =>
-        getTkmsPublicKeyHex(runtime, {
-          tkmsPrivateKey: privateKey,
-        }),
-      serializeTkmsPrivateKey: () =>
-        serializeTkmsPrivateKey(runtime, {
-          tkmsPrivateKey: privateKey,
-        }),
-    }),
-  });
-};
+// export const userDecryptModule: UserDecryptModuleFactory = (
+//   runtime: FhevmRuntime,
+//   parameters: UserDecryptModuleParameters,
+// ) => {
+//   const { privateKey } = parameters;
+//   return Object.freeze({
+//     userDecrypt: Object.freeze({
+//       initTkmsModule: async () => {
+//         await initTkmsModule(runtime);
+//       },
+//       getTkmsModuleInfo: () => getTkmsModuleInfo(),
+//       decryptAndReconstruct: (args: DecryptAndReconstructUserParameters) =>
+//         decryptAndReconstruct(runtime, {
+//           ...args,
+//           tkmsPrivateKey: privateKey,
+//         }),
+//       getTkmsPublicKeyHex: () =>
+//         getTkmsPublicKeyHex(runtime, {
+//           tkmsPrivateKey: privateKey,
+//         }),
+//       serializeTkmsPrivateKey: () =>
+//         serializeTkmsPrivateKey(runtime, {
+//           tkmsPrivateKey: privateKey,
+//         }),
+//     }),
+//   });
+// };

package/core/modules/decrypt/module/init-p.ts

@@ -1,16 +1,26 @@
 import type { FhevmRuntime, FhevmRuntimeConfig } from '../../../types/coreFhevmRuntime.js';
-import type { GetTkmsModuleInfoReturnType, TkmsModuleInfo } from '../types.js';
+import type {
+  GetTkmsModuleInfoParameters,
+  GetTkmsModuleInfoReturnType,
+  InitTkmsModuleParameters,
+  TkmsModuleInfo,
+} from '../types.js';
 import type { KmsLibApi, TkmsVersion } from '../../../../wasm/tkms/KmsLibApi.js';
 import type { KmsAssetMetadata } from '../../../../wasm/tkms/loadKmsLib.js';
 import { isBrowserLike } from '../../../base/isomorphicWorker.js';
 import { isomorphicCompileVerifiedWasm, isomorphicCompileWasmFromBase64 } from '../../../base/wasm.js';
 import { assertIsFhevmRuntime } from '../../../runtime/CoreFhevmRuntime-p.js';
-import { resolveTkmsModuleVersion } from '../../../../wasm/tkmsModuleVersion-p.js';
 import { wasmBaseUrl } from '../../../../wasm/wasmBaseUrl.js';
 import { kmsAssetsWithVersion, loadKmsLib, loadKmsWasmBase64 } from '../../../../wasm/tkms/loadKmsLib.js';
 
 ////////////////////////////////////////////////////////////////////////////////
 
+const ownerUidByVersion = new Map<TkmsVersion, string>();
+const cachedTkmsModulePromiseByVersion = new Map<TkmsVersion, Promise<KmsLibApi>>();
+const moduleInfoByVersion = new Map<TkmsVersion, TkmsModuleInfo>();
+
+////////////////////////////////////////////////////////////////////////////////
+
 // Pure JS file (not compiled by tsc) — provides cross-platform base URL
 // for resolving WASM paths. Uses import.meta.url in ESM, __filename in CJS.
 
@@ -41,26 +51,26 @@ type ResolvedTkmsModuleConfig = {
   readonly logger: FhevmRuntimeConfig['logger'];
 };
 
-let resolvedTkmsModuleConfig: ResolvedTkmsModuleConfig | undefined = undefined;
+const resolvedTkmsModuleConfigByVersion = new Map<TkmsVersion, ResolvedTkmsModuleConfig>();
 
 /**
  * @internal
  * Returns the existing resolved config, or resolves it from the runtime config.
  */
-function _getOrResolveTkmsModuleConfig(runtime: FhevmRuntime): ResolvedTkmsModuleConfig {
-  if (resolvedTkmsModuleConfig !== undefined) return resolvedTkmsModuleConfig;
+function _getOrResolveTkmsModuleConfig(runtime: FhevmRuntime, tkmsVersion: TkmsVersion): ResolvedTkmsModuleConfig {
+  let resolvedTkmsModuleConfig = resolvedTkmsModuleConfigByVersion.get(tkmsVersion);
+  if (resolvedTkmsModuleConfig !== undefined) {
+    return resolvedTkmsModuleConfig;
+  }
 
-  const version = resolveTkmsModuleVersion(runtime.config);
+  resolvedTkmsModuleConfig = _resolveTkmsModuleConfig(runtime.config, tkmsVersion);
+
+  resolvedTkmsModuleConfigByVersion.set(tkmsVersion, resolvedTkmsModuleConfig);
 
-  resolvedTkmsModuleConfig = _resolveTkmsModuleConfig(runtime.config, version);
   return resolvedTkmsModuleConfig;
 }
 
 function _resolveTkmsModuleConfig(parameters: FhevmRuntimeConfig, version: TkmsVersion): ResolvedTkmsModuleConfig {
-  if (cachedTkmsModulePromise !== undefined) {
-    throw new Error('Cannot configure module after initialization has started');
-  }
-
   const { locateFile } = parameters;
 
   const assets = kmsAssetsWithVersion(version);
@@ -95,23 +105,21 @@ function _resolveTkmsModuleConfig(parameters: FhevmRuntimeConfig, version: TkmsV
 
 type TkmsInitInput = RequestInfo | URL | Response | BufferSource | WebAssembly.Module;
 
-type InitTkmsModuleParameters = {
+type TkmsLibInitAsyncParameters = {
   readonly module_or_path: TkmsInitInput | Promise<TkmsInitInput>;
 };
 
-let cachedTkmsModulePromise: Promise<KmsLibApi> | undefined;
-let ownerUid: string | undefined = undefined;
-
-export async function initTkmsModule(runtime: FhevmRuntime): Promise<KmsLibApi> {
+export async function initTkmsModule(runtime: FhevmRuntime, parameters: InitTkmsModuleParameters): Promise<KmsLibApi> {
   assertIsFhevmRuntime(runtime, {});
 
+  const ownerUid = ownerUidByVersion.get(parameters.tkmsVersion);
   if (ownerUid !== undefined && runtime.uid !== ownerUid) {
     throw new Error(
       `Decrypt WASM module is already owned by runtime '${ownerUid}' and cannot be shared with runtime '${runtime.uid}'`,
     );
   }
 
-  ownerUid = runtime.uid;
+  ownerUidByVersion.set(parameters.tkmsVersion, runtime.uid);
 
   // Cache the TKMS initialization promise once and keep it cached even if it
   // rejects.
@@ -122,17 +130,18 @@ export async function initTkmsModule(runtime: FhevmRuntime): Promise<KmsLibApi>
   // keep module-level state that cannot be reset reliably after a partial
   // initialization failure; later callers should observe the original error
   // instead of retrying against half-initialized state.
+
+  let cachedTkmsModulePromise = cachedTkmsModulePromiseByVersion.get(parameters.tkmsVersion);
   if (cachedTkmsModulePromise === undefined) {
     // resolve is sync
-    const cfg = _getOrResolveTkmsModuleConfig(runtime);
+    const cfg = _getOrResolveTkmsModuleConfig(runtime, parameters.tkmsVersion);
     cachedTkmsModulePromise = _initTkmsModule(cfg);
+    cachedTkmsModulePromiseByVersion.set(parameters.tkmsVersion, cachedTkmsModulePromise);
   }
 
   return cachedTkmsModulePromise;
 }
 
-let moduleInfo: TkmsModuleInfo | undefined = undefined;
-
 async function _initTkmsModule(cfg: ResolvedTkmsModuleConfig): Promise<KmsLibApi> {
   const kmsLib = await loadKmsLib(cfg.version);
 
@@ -147,18 +156,27 @@ async function _initTkmsModule(cfg: ResolvedTkmsModuleConfig): Promise<KmsLibApi
     wasmModule = await isomorphicCompileWasmFromBase64(tkmsWasmBase64, tkmsWasmBase64IsGzipped ? 'gzip' : undefined);
   }
 
-  const input: InitTkmsModuleParameters = { module_or_path: wasmModule };
+  const input: TkmsLibInitAsyncParameters = { module_or_path: wasmModule };
 
   await kmsLib.initAsync(input);
 
   // Note: init_panic_hook is not exposed by kms_lib
   const wasmInfo = kmsLib.getWasmInfo();
+  const memory = { byteLength: 0, pages: 0 };
+  if (wasmInfo.memory !== undefined) {
+    memory.byteLength = wasmInfo.memory.byteLength;
+    memory.pages = wasmInfo.memory.pages;
+  }
 
-  moduleInfo = Object.freeze({
-    wasmUrl: cfg.wasm.url ? new URL(cfg.wasm.url) : undefined,
-    version: wasmInfo.version,
-    name: wasmInfo.name,
-  });
+  moduleInfoByVersion.set(
+    cfg.version,
+    Object.freeze({
+      wasmUrl: cfg.wasm.url ? new URL(cfg.wasm.url) : undefined,
+      version: wasmInfo.version,
+      name: wasmInfo.name,
+      memory,
+    }),
+  );
 
   return kmsLib;
 }
@@ -167,6 +185,20 @@ async function _initTkmsModule(cfg: ResolvedTkmsModuleConfig): Promise<KmsLibApi
 // getTkmsModuleInfo
 ////////////////////////////////////////////////////////////////////////////////
 
-export function getTkmsModuleInfo(): GetTkmsModuleInfoReturnType {
-  return moduleInfo;
+export async function getTkmsModuleInfo(parameters: GetTkmsModuleInfoParameters): Promise<GetTkmsModuleInfoReturnType> {
+  const tkmsLib = await loadKmsLib(parameters.tkmsVersion);
+  const stored = moduleInfoByVersion.get(parameters.tkmsVersion);
+  if (stored === undefined) {
+    throw new Error(`getTkmsModuleInfo: no module info recorded for version "${parameters.tkmsVersion}"`);
+  }
+  const memory: { byteLength: number; pages: number } = {
+    byteLength: stored.memory.byteLength,
+    pages: stored.memory.pages,
+  };
+  const wasmInfo = tkmsLib.getWasmInfo();
+  if (wasmInfo.memory !== undefined) {
+    memory.byteLength = wasmInfo.memory.byteLength;
+    memory.pages = wasmInfo.memory.pages;
+  }
+  return { ...stored, memory };
 }