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

package/core/modules/decrypt/types.ts

@@ -4,6 +4,7 @@ import type { Bytes, BytesHex } from '../../types/primitives.js';
 import type { Prettify } from '../../types/utils.js';
 import type { FhevmRuntime } from '../../types/coreFhevmRuntime.js';
 import type { ClearValue } from '../../types/encryptedTypes-p.js';
+import type { TkmsVersion } from '../../../wasm/tkms/KmsLibApi.js';
 
 ////////////////////////////////////////////////////////////////////////////////
 //
@@ -22,12 +23,20 @@ WASM compilation (how to get WebAssembly.Module):
 
 */
 
+type WithTkmsVersion = {
+  readonly tkmsVersion: TkmsVersion;
+};
+
 ////////////////////////////////////////////////////////////////////////////////
 // initTkmsModule
 ////////////////////////////////////////////////////////////////////////////////
 
+export type InitTkmsModuleParameters = {
+  readonly tkmsVersion: TkmsVersion;
+};
+
 export type InitTkmsModuleFunction = {
-  initTkmsModule(): Promise<void>;
+  initTkmsModule(parameters: InitTkmsModuleParameters): Promise<void>;
 };
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -39,12 +48,29 @@ export type InitTkmsModuleFunction = {
  */
 export type TkmsModuleInfo = {
   /**
-   * URL used to fetch the TFHE WASM binary.
+   * URL used to fetch the TKMS WASM binary.
    * `undefined` means the base64-embedded fallback is used.
    */
   readonly wasmUrl: URL | undefined;
+
+  /**
+   * Current size of the TKMS WASM linear memory.
+   * This is the size of the underlying `WebAssembly.Memory` buffer, not live heap usage.
+   */
+  readonly memory: {
+    /**
+     * Current memory buffer size in bytes.
+     */
+    readonly byteLength: number;
+    /**
+     * Current memory buffer size in 64 KiB WASM pages.
+     */
+    readonly pages: number;
+  };
 };
 
+export type GetTkmsModuleInfoParameters = WithTkmsVersion;
+
 export type GetTkmsModuleInfoReturnType = TkmsModuleInfo | undefined;
 
 export type GetTkmsModuleInfoFunction = {
@@ -52,11 +78,12 @@ export type GetTkmsModuleInfoFunction = {
    * Returns {@link TkmsModuleInfo} when the module is initialized,
    * or `undefined` if the module has not completed initialization.
    */
-  getTkmsModuleInfo(): GetTkmsModuleInfoReturnType;
+  getTkmsModuleInfo(parameters: GetTkmsModuleInfoParameters): Promise<GetTkmsModuleInfoReturnType>;
 };
+
 ////////////////////////////////////////////////////////////////////////////////
 
-type WithTkmsPrivateKey = {
+type WithTkmsPrivateKey = WithTkmsVersion & {
   readonly tkmsPrivateKey: TkmsPrivateKey;
 };
 
@@ -75,52 +102,33 @@ export type DecryptAndReconstructModuleFunction = {
   decryptAndReconstruct(parameters: DecryptAndReconstructParameters): Promise<DecryptAndReconstructReturnType>;
 };
 
-////////////////////////////////////////////////////////////////////////////////
-// 1. (User) decryptAndReconstruct
-////////////////////////////////////////////////////////////////////////////////
-
-export type DecryptAndReconstructUserParameters = DecryptAndReconstructBaseParameters;
-
-export type DecryptAndReconstructUserModuleFunction = {
-  decryptAndReconstruct(parameters: DecryptAndReconstructUserParameters): Promise<DecryptAndReconstructReturnType>;
-};
-
 ////////////////////////////////////////////////////////////////////////////////
 // 2. generateTkmsPrivateKey
 ////////////////////////////////////////////////////////////////////////////////
 
+export type GenerateTkmsPrivateKeyParameters = WithTkmsVersion;
 export type GenerateTkmsPrivateKeyReturnType = TkmsPrivateKey;
 
 export type GenerateTkmsPrivateKeyModuleFunction = {
-  generateTkmsPrivateKey(): Promise<GenerateTkmsPrivateKeyReturnType>;
+  generateTkmsPrivateKey(parameters: GenerateTkmsPrivateKeyParameters): Promise<GenerateTkmsPrivateKeyReturnType>;
 };
 
 ////////////////////////////////////////////////////////////////////////////////
 // 3 getTkmsPublicKeyHex
 ////////////////////////////////////////////////////////////////////////////////
 
-export type GetTkmsPublicKeyHexParameters = WithTkmsPrivateKey;
+export type GetTkmsPublicKeyHexParameters = WithTkmsPrivateKey & WithTkmsVersion;
 export type GetTkmsPublicKeyHexReturnType = BytesHex;
 
 export type GetTkmsPublicKeyHexModuleFunction = {
   getTkmsPublicKeyHex(parameters: GetTkmsPublicKeyHexParameters): Promise<GetTkmsPublicKeyHexReturnType>;
 };
 
-////////////////////////////////////////////////////////////////////////////////
-// 3. (User) getTkmsPublicKeyHex
-////////////////////////////////////////////////////////////////////////////////
-
-export type GetTkmsPublicKeyHexUserModuleFunction = {
-  getTkmsPublicKeyHex(): Promise<GetTkmsPublicKeyHexReturnType>;
-};
-
 ////////////////////////////////////////////////////////////////////////////////
 // 4. serializeTkmsPrivateKey
 ////////////////////////////////////////////////////////////////////////////////
 
-export type SerializeTkmsPrivateKeyParameters = {
-  readonly tkmsPrivateKey: TkmsPrivateKey;
-};
+export type SerializeTkmsPrivateKeyParameters = WithTkmsPrivateKey;
 
 export type SerializeTkmsPrivateKeyReturnType = Bytes;
 
@@ -128,19 +136,11 @@ export type SerializeTkmsPrivateKeyModuleFunction = {
   serializeTkmsPrivateKey(parameters: SerializeTkmsPrivateKeyParameters): Promise<SerializeTkmsPrivateKeyReturnType>;
 };
 
-////////////////////////////////////////////////////////////////////////////////
-// 4. (User) serializeTkmsPrivateKey
-////////////////////////////////////////////////////////////////////////////////
-
-export type SerializeTkmsPrivateKeyUserModuleFunction = {
-  serializeTkmsPrivateKey(): Promise<SerializeTkmsPrivateKeyReturnType>;
-};
-
 ////////////////////////////////////////////////////////////////////////////////
 // 5. deserializeTkmsPrivateKey
 ////////////////////////////////////////////////////////////////////////////////
 
-export type DeserializeTkmsPrivateKeyParameters = {
+export type DeserializeTkmsPrivateKeyParameters = WithTkmsVersion & {
   readonly tkmsPrivateKeyBytes: Bytes;
 };
 
@@ -156,9 +156,7 @@ export type DeserializeTkmsPrivateKeyModuleFunction = {
 // 6. verifyTkmsPrivateKey
 ////////////////////////////////////////////////////////////////////////////////
 
-export type VerifyTkmsPrivateKeyParameters = {
-  readonly tkmsPrivateKey: TkmsPrivateKey;
-};
+export type VerifyTkmsPrivateKeyParameters = WithTkmsVersion & WithTkmsPrivateKey;
 
 export type VerifyTkmsPrivateKeyModuleFunction = {
   verifyTkmsPrivateKey(parameters: VerifyTkmsPrivateKeyParameters): void;
@@ -186,29 +184,3 @@ export type DecryptModule = Prettify<
 export type DecryptModuleFactory = (runtime: FhevmRuntime) => {
   readonly decrypt: DecryptModule;
 };
-
-////////////////////////////////////////////////////////////////////////////////
-// UserDecryptModule
-////////////////////////////////////////////////////////////////////////////////
-
-export type WithUserDecryptModule = {
-  readonly userDecrypt: UserDecryptModule;
-};
-
-export type UserDecryptModuleParameters = {
-  readonly privateKey: TkmsPrivateKey;
-};
-
-export type UserDecryptModule = Prettify<
-  InitTkmsModuleFunction &
-    DecryptAndReconstructUserModuleFunction &
-    GetTkmsPublicKeyHexUserModuleFunction &
-    SerializeTkmsPrivateKeyUserModuleFunction
->;
-
-export type UserDecryptModuleFactory = (
-  runtime: FhevmRuntime,
-  parameters: UserDecryptModuleParameters,
-) => {
-  readonly userDecrypt: UserDecryptModule;
-};

package/core/modules/encrypt/mock.ts

@@ -1,4 +1,5 @@
 /* eslint-disable @typescript-eslint/require-await */
+import type { TfheVersion } from '../../../wasm/tfhe/TfheApi.js';
 import { concatBytes, hexToBytes32 } from '../../base/bytes.js';
 import { remove0x } from '../../base/string.js';
 import { typedValueToBytes32Hex } from '../../base/typedValue.js';
@@ -43,12 +44,14 @@ class CleartextTfheCompactPublicKeyImpl implements FheEncryptionPublicKey {
 
   readonly #id: string;
   readonly #tfheCompactPublicKeyWasmType: Bytes;
+  readonly #tfheVersion: TfheVersion;
 
-  constructor(token: symbol, id: string, publicEncKeyMlKem512Wasm: Bytes) {
+  constructor(token: symbol, id: string, tfheVersion: TfheVersion, publicEncKeyMlKem512Wasm: Bytes) {
     if (token !== PRIVATE_CLEARTEXT_TFHE_LIB_TOKEN) {
       throw new Error('Unauthorized');
     }
     this.#id = id;
+    this.#tfheVersion = tfheVersion;
     this.#tfheCompactPublicKeyWasmType = publicEncKeyMlKem512Wasm;
   }
 
@@ -56,6 +59,10 @@ class CleartextTfheCompactPublicKeyImpl implements FheEncryptionPublicKey {
     return this.#id;
   }
 
+  public get tfheVersion(): TfheVersion {
+    return this.#tfheVersion;
+  }
+
   public static [GET_NATIVE_FUNC](key: unknown, token: symbol): Bytes {
     if (token !== PRIVATE_CLEARTEXT_TFHE_LIB_TOKEN) {
       throw new Error('Unauthorized');
@@ -75,14 +82,22 @@ class CleartextTfheCompactPkeCrsImpl implements FheEncryptionCrs {
   declare readonly [FheEncryptionCrsBrand]: never;
 
   readonly #id: string;
+  readonly #tfheVersion: TfheVersion;
   readonly #capacity: UintNumber;
   readonly #compactPublicKeyWasmType: Bytes;
 
-  constructor(token: symbol, id: string, capacity: UintNumber, compactPublicKeyWasmType: Bytes) {
+  constructor(
+    token: symbol,
+    id: string,
+    tfheVersion: TfheVersion,
+    capacity: UintNumber,
+    compactPublicKeyWasmType: Bytes,
+  ) {
     if (token !== PRIVATE_CLEARTEXT_TFHE_LIB_TOKEN) {
       throw new Error('Unauthorized');
     }
     this.#id = id;
+    this.#tfheVersion = tfheVersion;
     this.#capacity = capacity;
     this.#compactPublicKeyWasmType = compactPublicKeyWasmType;
   }
@@ -91,6 +106,10 @@ class CleartextTfheCompactPkeCrsImpl implements FheEncryptionCrs {
     return this.#id;
   }
 
+  public get tfheVersion(): TfheVersion {
+    return this.#tfheVersion;
+  }
+
   public get capacity(): UintNumber {
     return this.#capacity;
   }
@@ -150,6 +169,7 @@ export async function buildWithProofPacked(
   return {
     ciphertextWithZKProofBytes,
     extraData: cleartextExtraData,
+    tfheVersion: parameters.tfheVersion,
   };
 }
 
@@ -181,7 +201,7 @@ export async function serializeFheEncryptionKey(
     PRIVATE_CLEARTEXT_TFHE_LIB_TOKEN,
   );
 
-  return Object.freeze({
+  const keyBytes: FheEncryptionKeyBytes = Object.freeze({
     publicKeyBytes: Object.freeze({
       id: publicEncryptionParams.publicKey.id,
       bytes: tfhePublicKeyBytes,
@@ -191,7 +211,13 @@ export async function serializeFheEncryptionKey(
       capacity: publicEncryptionParams.crs.capacity,
       bytes: tfheCrsBytes,
     }),
-  }) as FheEncryptionKeyBytes;
+    metadata: Object.freeze({
+      relayerUrl: publicEncryptionParams.metadata.relayerUrl,
+      chainId: publicEncryptionParams.metadata.chainId,
+    }),
+  });
+
+  return keyBytes;
 }
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -217,6 +243,7 @@ export async function serializeFheEncryptionPublicKey(
   return Object.freeze({
     id: tfhePublicKey.id,
     bytes: tfhePublicKeyBytes,
+    tfheVersion: parameters.tfheVersion,
   });
 }
 
@@ -244,6 +271,7 @@ export async function serializeFheEncryptionCrs(
     id: tfheCrs.id,
     capacity: tfheCrs.capacity,
     bytes: tfheCrsBytes,
+    tfheVersion: parameters.tfheVersion,
   });
 }
 
@@ -259,6 +287,7 @@ export async function deserializeFheEncryptionCrs(
   return new CleartextTfheCompactPkeCrsImpl(
     PRIVATE_CLEARTEXT_TFHE_LIB_TOKEN,
     globalFheCrsBytes.id,
+    parameters.tfheVersion,
     globalFheCrsBytes.capacity,
     globalFheCrsBytes.bytes,
   );
@@ -276,6 +305,7 @@ export async function deserializeFheEncryptionPublicKey(
   return new CleartextTfheCompactPublicKeyImpl(
     PRIVATE_CLEARTEXT_TFHE_LIB_TOKEN,
     globalFhePublicKeyBytes.id,
+    parameters.tfheVersion,
     globalFhePublicKeyBytes.bytes,
   );
 }

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

@@ -29,6 +29,7 @@ import type {
   TfheCompactPublicKey,
   ProvenCompactCiphertextList,
   CompactCiphertextListBuilder,
+  TfheVersion,
 } from '../../../../wasm/tfhe/TfheApi.js';
 import { isNonEmptyString } from '../../../base/string.js';
 import { hexToBytesFaster } from '../../../base/bytes.js';
@@ -62,12 +63,14 @@ class TfheCompactPublicKeyImpl implements FheEncryptionPublicKey {
 
   readonly #id: string;
   readonly #tfheCompactPublicKeyWasmType: TfheCompactPublicKey;
+  readonly #tfheVersion: TfheVersion;
 
-  constructor(token: symbol, id: string, publicEncKeyMlKem512Wasm: TfheCompactPublicKey) {
+  constructor(token: symbol, id: string, tfheVersion: TfheVersion, publicEncKeyMlKem512Wasm: TfheCompactPublicKey) {
     if (token !== PRIVATE_TFHE_LIB_TOKEN) {
       throw new Error('Unauthorized');
     }
     this.#id = id;
+    this.#tfheVersion = tfheVersion;
     this.#tfheCompactPublicKeyWasmType = publicEncKeyMlKem512Wasm;
   }
 
@@ -75,13 +78,20 @@ class TfheCompactPublicKeyImpl implements FheEncryptionPublicKey {
     return this.#id;
   }
 
-  public static [GET_NATIVE_FUNC](key: unknown, token: symbol): TfheCompactPublicKey {
+  public get tfheVersion(): TfheVersion {
+    return this.#tfheVersion;
+  }
+
+  public static [GET_NATIVE_FUNC](key: unknown, token: symbol, tfheVersion: TfheVersion): TfheCompactPublicKey {
     if (token !== PRIVATE_TFHE_LIB_TOKEN) {
       throw new Error('Unauthorized');
     }
     if (!(key instanceof TfheCompactPublicKeyImpl)) {
       throw new Error('Unauthorized');
     }
+    if (tfheVersion !== key.#tfheVersion) {
+      throw new Error('TfheVersion mismatch');
+    }
     return key.#tfheCompactPublicKeyWasmType;
   }
 }
@@ -96,12 +106,20 @@ class TfheCompactPkeCrsImpl implements FheEncryptionCrs {
   readonly #id: string;
   readonly #capacity: UintNumber;
   readonly #compactPublicKeyWasmType: CompactPkeCrs;
-
-  constructor(token: symbol, id: string, capacity: UintNumber, compactPublicKeyWasmType: CompactPkeCrs) {
+  readonly #tfheVersion: TfheVersion;
+
+  constructor(
+    token: symbol,
+    id: string,
+    tfheVersion: TfheVersion,
+    capacity: UintNumber,
+    compactPublicKeyWasmType: CompactPkeCrs,
+  ) {
     if (token !== PRIVATE_TFHE_LIB_TOKEN) {
       throw new Error('Unauthorized');
     }
     this.#id = id;
+    this.#tfheVersion = tfheVersion;
     this.#capacity = capacity;
     this.#compactPublicKeyWasmType = compactPublicKeyWasmType;
   }
@@ -110,17 +128,24 @@ class TfheCompactPkeCrsImpl implements FheEncryptionCrs {
     return this.#id;
   }
 
+  public get tfheVersion(): TfheVersion {
+    return this.#tfheVersion;
+  }
+
   public get capacity(): UintNumber {
     return this.#capacity;
   }
 
-  public static [GET_NATIVE_FUNC](key: unknown, token: symbol): CompactPkeCrs {
+  public static [GET_NATIVE_FUNC](key: unknown, token: symbol, tfheVersion: TfheVersion): CompactPkeCrs {
     if (token !== PRIVATE_TFHE_LIB_TOKEN) {
       throw new Error('Unauthorized');
     }
     if (!(key instanceof TfheCompactPkeCrsImpl)) {
       throw new Error('Unauthorized');
     }
+    if (tfheVersion !== key.#tfheVersion) {
+      throw new Error('TfheVersion mismatch');
+    }
     return key.#compactPublicKeyWasmType;
   }
 }
@@ -133,7 +158,7 @@ export async function parseTFHEProvenCompactCiphertextList(
   runtime: FhevmRuntime,
   parameters: ParseTFHEProvenCompactCiphertextListParameters,
 ): Promise<ParseTFHEProvenCompactCiphertextListReturnType> {
-  const tfheLib = await initTfheModule(runtime);
+  const tfheLib = await initTfheModule(runtime, { tfheVersion: parameters.tfheVersion });
 
   const { ciphertextWithZkProof: ciphertextWithZKProof } = parameters;
   if ((ciphertextWithZKProof as unknown) == null) {
@@ -193,7 +218,7 @@ export async function buildWithProofPacked(
   runtime: FhevmRuntime,
   parameters: BuildWithProofPackedParameters,
 ): Promise<BuildWithProofPackedReturnType> {
-  const tfheLib = await initTfheModule(runtime);
+  const tfheLib = await initTfheModule(runtime, { tfheVersion: parameters.tfheVersion });
 
   const { fheEncryptionKey: publicEncryptionParams, metaData, typedValues, extraData } = parameters;
 
@@ -216,10 +241,12 @@ export async function buildWithProofPacked(
     const tfheCompactPublicKeyWasm: TfheCompactPublicKey = TfheCompactPublicKeyImpl[GET_NATIVE_FUNC](
       tfheCompactPublicKeyImpl,
       PRIVATE_TFHE_LIB_TOKEN,
+      parameters.tfheVersion,
     );
     const compactPkeCrsWasm: CompactPkeCrs = TfheCompactPkeCrsImpl[GET_NATIVE_FUNC](
       tfheCompactPkeCrsImpl,
       PRIVATE_TFHE_LIB_TOKEN,
+      parameters.tfheVersion,
     );
 
     fheCompactCiphertextListBuilderWasm = tfheLib.CompactCiphertextList.builder(tfheCompactPublicKeyWasm);
@@ -265,6 +292,7 @@ export async function buildWithProofPacked(
     return Object.freeze({
       ciphertextWithZKProofBytes,
       extraData,
+      tfheVersion: parameters.tfheVersion,
     });
   } finally {
     try {
@@ -293,7 +321,8 @@ export async function serializeFheEncryptionKey(
   runtime: FhevmRuntime,
   parameters: SerializeFheEncryptionKeyParameters,
 ): Promise<SerializeFheEncryptionKeyReturnType> {
-  await initTfheModule(runtime);
+  const tfheVersion = parameters.tfheVersion;
+  await initTfheModule(runtime, { tfheVersion });
 
   const { fheEncryptionKey: publicEncryptionParams } = parameters;
 
@@ -310,13 +339,15 @@ export async function serializeFheEncryptionKey(
   const tfhePublicKeyBytes: Bytes = TfheCompactPublicKeyImpl[GET_NATIVE_FUNC](
     tfheCompactPublicKeyImpl,
     PRIVATE_TFHE_LIB_TOKEN,
+    tfheVersion,
   ).safe_serialize(SERIALIZED_SIZE_LIMIT_PK);
   const tfheCrsBytes: Bytes = TfheCompactPkeCrsImpl[GET_NATIVE_FUNC](
     tfheCompactPkeCrsImpl,
     PRIVATE_TFHE_LIB_TOKEN,
+    tfheVersion,
   ).safe_serialize(SERIALIZED_SIZE_LIMIT_CRS);
 
-  return Object.freeze({
+  const keyBytes: FheEncryptionKeyBytes = Object.freeze({
     publicKeyBytes: Object.freeze({
       id: publicEncryptionParams.publicKey.id,
       bytes: tfhePublicKeyBytes,
@@ -326,7 +357,13 @@ export async function serializeFheEncryptionKey(
       capacity: publicEncryptionParams.crs.capacity,
       bytes: tfheCrsBytes,
     }),
-  }) as FheEncryptionKeyBytes;
+    metadata: Object.freeze({
+      relayerUrl: publicEncryptionParams.metadata.relayerUrl,
+      chainId: publicEncryptionParams.metadata.chainId,
+    }),
+  });
+
+  return keyBytes;
 }
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -337,7 +374,8 @@ export async function serializeFheEncryptionPublicKey(
   runtime: FhevmRuntime,
   parameters: SerializeFheEncryptionPublicKeyParameters,
 ): Promise<SerializeFheEncryptionPublicKeyReturnType> {
-  await initTfheModule(runtime);
+  const tfheVersion = parameters.tfheVersion;
+  await initTfheModule(runtime, { tfheVersion });
 
   const { publicKey: tfhePublicKey } = parameters;
 
@@ -350,11 +388,13 @@ export async function serializeFheEncryptionPublicKey(
   const tfhePublicKeyBytes: Bytes = TfheCompactPublicKeyImpl[GET_NATIVE_FUNC](
     tfheCompactPublicKeyImpl,
     PRIVATE_TFHE_LIB_TOKEN,
+    tfheVersion,
   ).safe_serialize(SERIALIZED_SIZE_LIMIT_PK);
 
   return Object.freeze({
     id: tfhePublicKey.id,
     bytes: tfhePublicKeyBytes,
+    tfheVersion,
   });
 }
 
@@ -366,7 +406,8 @@ export async function serializeFheEncryptionCrs(
   runtime: FhevmRuntime,
   parameters: SerializeFheEncryptionCrsParameters,
 ): Promise<SerializeFheEncryptionCrsReturnType> {
-  await initTfheModule(runtime);
+  const tfheVersion = parameters.tfheVersion;
+  await initTfheModule(runtime, { tfheVersion });
 
   const { crs: tfheCrs } = parameters;
 
@@ -379,12 +420,14 @@ export async function serializeFheEncryptionCrs(
   const tfheCrsBytes: Bytes = TfheCompactPkeCrsImpl[GET_NATIVE_FUNC](
     tfheCompactPkeCrsImpl,
     PRIVATE_TFHE_LIB_TOKEN,
+    tfheVersion,
   ).safe_serialize(SERIALIZED_SIZE_LIMIT_CRS);
 
   return Object.freeze({
     id: tfheCrs.id,
     capacity: tfheCrs.capacity,
     bytes: tfheCrsBytes,
+    tfheVersion,
   });
 }
 
@@ -396,7 +439,8 @@ export async function deserializeFheEncryptionCrs(
   runtime: FhevmRuntime,
   parameters: DeserializeFheEncryptionCrsParameters,
 ): Promise<DeserializeFheEncryptionCrsReturnType> {
-  const tfheLib = await initTfheModule(runtime);
+  const tfheVersion = parameters.tfheVersion;
+  const tfheLib = await initTfheModule(runtime, { tfheVersion });
 
   const { crsBytes: globalFheCrsBytes } = parameters;
 
@@ -408,6 +452,7 @@ export async function deserializeFheEncryptionCrs(
   return new TfheCompactPkeCrsImpl(
     PRIVATE_TFHE_LIB_TOKEN,
     globalFheCrsBytes.id,
+    tfheVersion,
     globalFheCrsBytes.capacity,
     compactPkeCrsWasm,
   );
@@ -421,7 +466,8 @@ export async function deserializeFheEncryptionPublicKey(
   runtime: FhevmRuntime,
   parameters: DeserializeFheEncryptionPublicKeyParameters,
 ): Promise<DeserializeFheEncryptionPublicKeyReturnType> {
-  const tfheLib = await initTfheModule(runtime);
+  const tfheVersion = parameters.tfheVersion;
+  const tfheLib = await initTfheModule(runtime, { tfheVersion });
 
   const { publicKeyBytes: globalFhePublicKeyBytes } = parameters;
 
@@ -430,5 +476,10 @@ export async function deserializeFheEncryptionPublicKey(
     SERIALIZED_SIZE_LIMIT_PK,
   );
 
-  return new TfheCompactPublicKeyImpl(PRIVATE_TFHE_LIB_TOKEN, globalFhePublicKeyBytes.id, tfheCompactPublicKeyWasm);
+  return new TfheCompactPublicKeyImpl(
+    PRIVATE_TFHE_LIB_TOKEN,
+    globalFhePublicKeyBytes.id,
+    tfheVersion,
+    tfheCompactPublicKeyWasm,
+  );
 }

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

@@ -4,6 +4,8 @@ import type {
   DeserializeFheEncryptionCrsParameters,
   DeserializeFheEncryptionPublicKeyParameters,
   EncryptModuleFactory,
+  GetTfheModuleInfoParameters,
+  InitTfheModuleParameters,
   ParseTFHEProvenCompactCiphertextListParameters,
   SerializeFheEncryptionCrsParameters,
   SerializeFheEncryptionKeyParameters,
@@ -27,10 +29,10 @@ import { getTfheModuleInfo, initTfheModule } from './init-p.js';
 export const encryptModule: EncryptModuleFactory = (runtime: FhevmRuntime) => {
   return Object.freeze({
     encrypt: Object.freeze({
-      initTfheModule: async () => {
-        await initTfheModule(runtime);
+      initTfheModule: async (args: InitTfheModuleParameters) => {
+        await initTfheModule(runtime, args);
       },
-      getTfheModuleInfo: () => getTfheModuleInfo(),
+      getTfheModuleInfo: (args: GetTfheModuleInfoParameters) => getTfheModuleInfo(args),
       parseTFHEProvenCompactCiphertextList: (args: ParseTFHEProvenCompactCiphertextListParameters) =>
         parseTFHEProvenCompactCiphertextList(runtime, args),
       buildWithProofPacked: (args: BuildWithProofPackedParameters) => buildWithProofPacked(runtime, args),