@cofhe/sdk 0.4.0 → 0.5.1

package/core/decrypt/tnSealOutputV2.ts

@@ -1,16 +1,38 @@
 import { type Permission, type EthEncryptedData } from '@/permits';
 
 import { CofheError, CofheErrorCode } from '../error.js';
+import { type DecryptPollCallbackFunction } from '../types.js';
+import { computeMinuteRampPollIntervalMs } from './polling.js';
 
 // Polling configuration
 const POLL_INTERVAL_MS = 1000; // 1 second
-const POLL_TIMEOUT_MS = 5 * 60 * 1000; // 5 minutes
+const POLL_MAX_INTERVAL_MS = 10_000; // 10 seconds
+const SEAL_OUTPUT_TIMEOUT_MS = 5 * 60 * 1000; // 5 minutes total across submit + poll
+const SUBMIT_RETRY_INTERVAL_MS = 1000; // 1 second
 
 // V2 API response types
 type SealOutputSubmitResponse = {
-  request_id: string;
+  request_id: string | null;
+  status?: string;
+  is_succeed?: boolean;
+  sealed?: {
+    data: number[];
+    public_key: number[];
+    nonce: number[];
+  };
+  sealed_data?: number[];
+  ephemeral_public_key?: number[];
+  nonce?: number[];
+  signature?: string;
+  encryption_type?: number;
+  error_message?: string | null;
+  message?: string;
 };
 
+type SealOutputSubmitResult =
+  | { kind: 'request_id'; requestId: string }
+  | { kind: 'completed'; sealed: EthEncryptedData };
+
 type SealOutputStatusResponse = {
   request_id: string;
   status: 'PROCESSING' | 'COMPLETED';
@@ -52,83 +74,192 @@ function convertSealedData(sealed: SealOutputStatusResponse['sealed']): EthEncry
   };
 }
 
-/**
- * Submits a sealoutput request to the v2 API and returns the request_id
- */
-async function submitSealOutputRequest(
-  thresholdNetworkUrl: string,
-  ctHash: bigint | string,
-  chainId: number,
-  permission: Permission
-): Promise<string> {
-  const body = {
-    ct_tempkey: BigInt(ctHash).toString(16).padStart(64, '0'),
-    host_chain_id: chainId,
-    permit: permission,
-  };
+function getSealedDataFromSubmitResponse(
+  value: SealOutputSubmitResponse
+): SealOutputStatusResponse['sealed'] | undefined {
+  if (value.sealed) return value.sealed;
 
-  let response: Response;
-  try {
-    response = await fetch(`${thresholdNetworkUrl}/v2/sealoutput`, {
-      method: 'POST',
-      headers: {
-        'Content-Type': 'application/json',
-      },
-      body: JSON.stringify(body),
-    });
-  } catch (e) {
-    throw new CofheError({
-      code: CofheErrorCode.SealOutputFailed,
-      message: `sealOutput request failed`,
-      hint: 'Ensure the threshold network URL is valid and reachable.',
-      cause: e instanceof Error ? e : undefined,
-      context: {
-        thresholdNetworkUrl,
-        body,
-      },
-    });
+  if (Array.isArray(value.sealed_data) && Array.isArray(value.ephemeral_public_key) && Array.isArray(value.nonce)) {
+    return {
+      data: value.sealed_data,
+      public_key: value.ephemeral_public_key,
+      nonce: value.nonce,
+    };
   }
 
-  // Handle non-200 status codes
-  if (!response.ok) {
-    let errorMessage = `HTTP ${response.status}`;
-    try {
-      const errorBody = await response.json();
-      errorMessage = errorBody.error_message || errorBody.message || errorMessage;
-    } catch {
-      // Ignore JSON parse errors, use status text
-      errorMessage = response.statusText || errorMessage;
-    }
+  return undefined;
+}
 
+function parseCompletedSealOutputResponse(params: {
+  value: Pick<SealOutputStatusResponse, 'sealed' | 'error_message' | 'is_succeed'>;
+  thresholdNetworkUrl: string;
+  requestId?: string | null;
+}): EthEncryptedData {
+  const { value, thresholdNetworkUrl, requestId } = params;
+
+  if (value.is_succeed === false) {
+    const errorMessage = value.error_message || 'Unknown error';
     throw new CofheError({
       code: CofheErrorCode.SealOutputFailed,
       message: `sealOutput request failed: ${errorMessage}`,
-      hint: 'Check the threshold network URL and request parameters.',
       context: {
         thresholdNetworkUrl,
-        status: response.status,
-        statusText: response.statusText,
-        body,
+        requestId,
+        response: value,
       },
     });
   }
 
-  let submitResponse: SealOutputSubmitResponse;
-  try {
-    submitResponse = (await response.json()) as SealOutputSubmitResponse;
-  } catch (e) {
+  const sealed = 'sealed' in value ? value.sealed : getSealedDataFromSubmitResponse(value as SealOutputSubmitResponse);
+
+  if (!sealed) {
     throw new CofheError({
-      code: CofheErrorCode.SealOutputFailed,
-      message: `Failed to parse sealOutput submit response`,
-      cause: e instanceof Error ? e : undefined,
+      code: CofheErrorCode.SealOutputReturnedNull,
+      message: `sealOutput request completed but returned no sealed data`,
       context: {
         thresholdNetworkUrl,
-        body,
+        requestId,
+        response: value,
       },
     });
   }
 
-  if (!submitResponse.request_id) {
+  return convertSealedData(sealed);
+}
+
+/**
+ * Submits a sealoutput request to the v2 API and returns the request_id
+ */
+async function submitSealOutputRequest(
+  thresholdNetworkUrl: string,
+  ctHash: bigint | string,
+  chainId: number,
+  permission: Permission,
+  overallStartTime: number,
+  onPoll?: DecryptPollCallbackFunction
+): Promise<SealOutputSubmitResult> {
+  const body = {
+    ct_tempkey: BigInt(ctHash).toString(16).padStart(64, '0'),
+    host_chain_id: chainId,
+    permit: permission,
+  };
+  let attemptIndex = 0;
+
+  for (;;) {
+    let response: Response;
+    try {
+      response = await fetch(`${thresholdNetworkUrl}/v2/sealoutput`, {
+        method: 'POST',
+        headers: {
+          'Content-Type': 'application/json',
+        },
+        body: JSON.stringify(body),
+      });
+    } catch (e) {
+      throw new CofheError({
+        code: CofheErrorCode.SealOutputFailed,
+        message: `sealOutput request failed`,
+        hint: 'Ensure the threshold network URL is valid and reachable.',
+        cause: e instanceof Error ? e : undefined,
+        context: {
+          thresholdNetworkUrl,
+          body,
+          attemptIndex,
+        },
+      });
+    }
+
+    // Handle non-200 status codes
+    if (!response.ok) {
+      let errorMessage = `HTTP ${response.status}`;
+      try {
+        const errorBody = await response.json();
+
+        errorMessage = errorBody.error_message || errorBody.message || errorMessage;
+      } catch {
+        errorMessage = response.statusText || errorMessage;
+      }
+
+      throw new CofheError({
+        code: CofheErrorCode.SealOutputFailed,
+        message: `sealOutput request failed: ${errorMessage}`,
+        hint: 'Check the threshold network URL and request parameters.',
+        context: {
+          thresholdNetworkUrl,
+          status: response.status,
+          statusText: response.statusText,
+          body,
+          attemptIndex,
+        },
+      });
+    }
+
+    let submitResponse: SealOutputSubmitResponse | undefined;
+    if (response.status !== 204) {
+      try {
+        submitResponse = (await response.json()) as SealOutputSubmitResponse;
+      } catch (e) {
+        throw new CofheError({
+          code: CofheErrorCode.SealOutputFailed,
+          message: `Failed to parse sealOutput submit response`,
+          cause: e instanceof Error ? e : undefined,
+          context: {
+            thresholdNetworkUrl,
+            body,
+            attemptIndex,
+          },
+        });
+      }
+
+      if (getSealedDataFromSubmitResponse(submitResponse)) {
+        return {
+          kind: 'completed',
+          sealed: parseCompletedSealOutputResponse({
+            value: submitResponse,
+            thresholdNetworkUrl,
+            requestId: submitResponse.request_id,
+          }),
+        };
+      }
+
+      if (submitResponse.request_id) {
+        return { kind: 'request_id', requestId: submitResponse.request_id };
+      }
+    }
+
+    // 204 means backend is aware of ct hash but didn't calculate it yet
+    if (response.status === 204) {
+      const elapsedMs = Date.now() - overallStartTime;
+      if (elapsedMs > SEAL_OUTPUT_TIMEOUT_MS) {
+        throw new CofheError({
+          code: CofheErrorCode.SealOutputFailed,
+          message: `sealOutput submit retried without receiving request_id for ${SEAL_OUTPUT_TIMEOUT_MS}ms`,
+          hint: 'The ciphertext may still be propagating. Try again later.',
+          context: {
+            thresholdNetworkUrl,
+            body,
+            attemptIndex,
+            timeoutMs: SEAL_OUTPUT_TIMEOUT_MS,
+            submitResponse,
+            status: response.status,
+          },
+        });
+      }
+
+      onPoll?.({
+        operation: 'sealoutput',
+        requestId: '',
+        attemptIndex,
+        elapsedMs,
+        intervalMs: SUBMIT_RETRY_INTERVAL_MS,
+        timeoutMs: SEAL_OUTPUT_TIMEOUT_MS,
+      });
+
+      await new Promise((resolve) => setTimeout(resolve, SUBMIT_RETRY_INTERVAL_MS));
+      attemptIndex += 1;
+      continue;
+    }
+
     throw new CofheError({
       code: CofheErrorCode.SealOutputFailed,
       message: `sealOutput submit response missing request_id`,
@@ -136,31 +267,49 @@ async function submitSealOutputRequest(
         thresholdNetworkUrl,
         body,
         submitResponse,
+        attemptIndex,
       },
     });
   }
-
-  return submitResponse.request_id;
 }
 
 /**
  * Polls for the sealoutput status until completed or timeout
  */
-async function pollSealOutputStatus(thresholdNetworkUrl: string, requestId: string): Promise<EthEncryptedData> {
-  const startTime = Date.now();
+async function pollSealOutputStatus(
+  thresholdNetworkUrl: string,
+  requestId: string,
+  overallStartTime: number,
+  onPoll?: DecryptPollCallbackFunction
+): Promise<EthEncryptedData> {
+  let attemptIndex = 0;
   let completed = false;
 
   while (!completed) {
+    const elapsedMs = Date.now() - overallStartTime;
+    const intervalMs = computeMinuteRampPollIntervalMs(elapsedMs, {
+      minIntervalMs: POLL_INTERVAL_MS,
+      maxIntervalMs: POLL_MAX_INTERVAL_MS,
+    });
+    onPoll?.({
+      operation: 'sealoutput',
+      requestId,
+      attemptIndex,
+      elapsedMs,
+      intervalMs,
+      timeoutMs: SEAL_OUTPUT_TIMEOUT_MS,
+    });
+
     // Check timeout
-    if (Date.now() - startTime > POLL_TIMEOUT_MS) {
+    if (elapsedMs > SEAL_OUTPUT_TIMEOUT_MS) {
       throw new CofheError({
         code: CofheErrorCode.SealOutputFailed,
-        message: `sealOutput polling timed out after ${POLL_TIMEOUT_MS}ms`,
+        message: `sealOutput polling timed out after ${SEAL_OUTPUT_TIMEOUT_MS}ms`,
         hint: 'The request may still be processing. Try again later.',
         context: {
           thresholdNetworkUrl,
           requestId,
-          timeoutMs: POLL_TIMEOUT_MS,
+          timeoutMs: SEAL_OUTPUT_TIMEOUT_MS,
         },
       });
     }
@@ -238,39 +387,16 @@ async function pollSealOutputStatus(thresholdNetworkUrl: string, requestId: stri
 
     // Check if completed
     if (statusResponse.status === 'COMPLETED') {
-      // Check if succeeded
-      if (statusResponse.is_succeed === false) {
-        const errorMessage = statusResponse.error_message || 'Unknown error';
-        throw new CofheError({
-          code: CofheErrorCode.SealOutputFailed,
-          message: `sealOutput request failed: ${errorMessage}`,
-          context: {
-            thresholdNetworkUrl,
-            requestId,
-            statusResponse,
-          },
-        });
-      }
-
-      // Check if sealed data exists
-      if (!statusResponse.sealed) {
-        throw new CofheError({
-          code: CofheErrorCode.SealOutputReturnedNull,
-          message: `sealOutput request completed but returned no sealed data`,
-          context: {
-            thresholdNetworkUrl,
-            requestId,
-            statusResponse,
-          },
-        });
-      }
-
-      // Convert and return the sealed data
-      return convertSealedData(statusResponse.sealed);
+      return parseCompletedSealOutputResponse({
+        value: statusResponse,
+        thresholdNetworkUrl,
+        requestId,
+      });
     }
 
     // Still processing, wait before next poll
-    await new Promise((resolve) => setTimeout(resolve, POLL_INTERVAL_MS));
+    await new Promise((resolve) => setTimeout(resolve, intervalMs));
+    attemptIndex += 1;
   }
 
   // This should never be reached, but TypeScript requires it
@@ -284,15 +410,30 @@ async function pollSealOutputStatus(thresholdNetworkUrl: string, requestId: stri
   });
 }
 
-export async function tnSealOutputV2(
-  ctHash: bigint | string,
-  chainId: number,
-  permission: Permission,
-  thresholdNetworkUrl: string
-): Promise<EthEncryptedData> {
+export async function tnSealOutputV2(params: {
+  ctHash: bigint | string;
+  chainId: number;
+  permission: Permission;
+  thresholdNetworkUrl: string;
+  onPoll?: DecryptPollCallbackFunction;
+}): Promise<EthEncryptedData> {
+  const { thresholdNetworkUrl, ctHash, chainId, permission, onPoll } = params;
+  const overallStartTime = Date.now();
+
   // Step 1: Submit the request and get request_id
-  const requestId = await submitSealOutputRequest(thresholdNetworkUrl, ctHash, chainId, permission);
+  const submitResult = await submitSealOutputRequest(
+    thresholdNetworkUrl,
+    ctHash,
+    chainId,
+    permission,
+    overallStartTime,
+    onPoll
+  );
+
+  if (submitResult.kind === 'completed') {
+    return submitResult.sealed;
+  }
 
   // Step 2: Poll for status until completed
-  return await pollSealOutputStatus(thresholdNetworkUrl, requestId);
+  return await pollSealOutputStatus(thresholdNetworkUrl, submitResult.requestId, overallStartTime, onPoll);
 }

package/core/decrypt/verifyDecryptResult.ts

@@ -0,0 +1,65 @@
+import {
+  encodePacked,
+  isAddressEqual,
+  keccak256,
+  parseAbi,
+  recoverAddress,
+  zeroAddress,
+  type Address,
+  type Hex,
+  type PublicClient,
+} from 'viem';
+import { TASK_MANAGER_ADDRESS } from '../consts.js';
+
+const decryptResultSignerAbi = parseAbi(['function decryptResultSigner() view returns (address)']);
+const UINT_TYPE_MASK = 0x7fn;
+const TYPE_BYTE_OFFSET = 8n;
+
+const getEncryptionTypeFromCtHash = (ctHash: bigint) => Number((ctHash >> TYPE_BYTE_OFFSET) & UINT_TYPE_MASK);
+
+const buildDecryptResultHash = (ctHash: bigint, cleartext: bigint, chainId: number) => {
+  const encryptionType = getEncryptionTypeFromCtHash(ctHash);
+
+  return keccak256(
+    encodePacked(['uint256', 'uint32', 'uint64', 'uint256'], [cleartext, encryptionType, BigInt(chainId), ctHash])
+  );
+};
+
+/**
+ * Verifies a decrypt result signature **locally** (no `ctHash`/plaintext sent over RPC).
+ *
+ * The recovered signer must equal the on-chain configured `decryptResultSigner`.
+ *
+ * This mirrors the TaskManager decrypt-result hash format:
+ * `keccak256(abi.encodePacked(result, encType, chainId, ctHash))`
+ *
+ * The only on-chain read performed is `TaskManager.decryptResultSigner()` (via `eth_call`).
+ *
+ * Works with both production and mock deployments.
+ */
+export async function verifyDecryptResult(
+  handle: bigint | string,
+  cleartext: bigint,
+  signature: Hex,
+  publicClient: PublicClient
+): Promise<boolean> {
+  const chainId = publicClient.chain?.id ?? (await publicClient.getChainId());
+  const expectedSigner = await publicClient.readContract({
+    address: TASK_MANAGER_ADDRESS,
+    abi: decryptResultSignerAbi,
+    functionName: 'decryptResultSigner',
+    args: [],
+  });
+
+  if (isAddressEqual(expectedSigner, zeroAddress)) return true;
+
+  const ctHash = BigInt(handle);
+  const messageHash = buildDecryptResultHash(ctHash, cleartext, chainId);
+
+  try {
+    const recovered = await recoverAddress({ hash: messageHash, signature });
+    return isAddressEqual(recovered, expectedSigner);
+  } catch {
+    return false;
+  }
+}

package/core/encrypt/cofheMocksZkVerifySign.ts

@@ -1,5 +1,5 @@
 import { type EncryptableItem, FheTypes } from '../types.js';
-import { MAX_ENCRYPTABLE_BITS, type VerifyResult } from './zkPackProveVerify.js';
+import { type VerifyResult } from './zkPackProveVerify.js';
 import {
   createWalletClient,
   http,
@@ -14,7 +14,7 @@ import { MockZkVerifierAbi } from './MockZkVerifierAbi.js';
 import { hardhat } from 'viem/chains';
 import { CofheError, CofheErrorCode } from '../error.js';
 import { privateKeyToAccount, sign } from 'viem/accounts';
-import { MOCKS_ZK_VERIFIER_SIGNER_PRIVATE_KEY, MOCKS_ZK_VERIFIER_ADDRESS } from '../consts.js';
+import { MOCKS_ZK_VERIFIER_SIGNER_PRIVATE_KEY, MOCKS_ZK_VERIFIER_ADDRESS, TFHE_RS_ZK_MAX_BITS } from '../consts.js';
 
 type EncryptableItemWithCtHash = EncryptableItem & {
   ctHash: bigint;
@@ -69,14 +69,14 @@ export async function cofheMocksCheckEncryptableBits(items: EncryptableItem[]):
       }
     }
   }
-  if (totalBits > MAX_ENCRYPTABLE_BITS) {
+  if (totalBits > TFHE_RS_ZK_MAX_BITS) {
     throw new CofheError({
       code: CofheErrorCode.ZkPackFailed,
-      message: `Total bits ${totalBits} exceeds ${MAX_ENCRYPTABLE_BITS}`,
-      hint: `Ensure that the total bits of the items to encrypt does not exceed ${MAX_ENCRYPTABLE_BITS}`,
+      message: `Total bits ${totalBits} exceeds ${TFHE_RS_ZK_MAX_BITS}`,
+      hint: `Ensure that the total bits of the items to encrypt does not exceed ${TFHE_RS_ZK_MAX_BITS}`,
       context: {
         totalBits,
-        maxBits: MAX_ENCRYPTABLE_BITS,
+        maxBits: TFHE_RS_ZK_MAX_BITS,
         items,
       },
     });

package/core/encrypt/zkPackProveVerify.ts

@@ -1,6 +1,7 @@
-import { CofheError, CofheErrorCode } from '../error.js';
-import { type EncryptableItem, FheTypes } from '../types.js';
-import { toBigIntOrThrow, validateBigIntInRange, toHexString, hexToBytes } from '../utils.js';
+import { TFHE_RS_ZK_MAX_BITS, TFHE_RS_SAFE_SERIALIZATION_SIZE_LIMIT } from '../consts';
+import { CofheError, CofheErrorCode } from '../error';
+import { type EncryptableItem, FheTypes } from '../types';
+import { toBigIntOrThrow, validateBigIntInRange, toHexString, hexToBytes } from '../utils';
 
 // ===== TYPES =====
 
@@ -52,23 +53,14 @@ export type VerifyResult = {
 };
 
 export type ZkProvenCiphertextList = {
-  serialize(): Uint8Array;
+  safe_serialize(serialized_size_limit: bigint): Uint8Array;
 };
 
 export type ZkCompactPkeCrs = {
   free(): void;
-  serialize(compress: boolean): Uint8Array;
   safe_serialize(serialized_size_limit: bigint): Uint8Array;
 };
 
-export type ZkCompactPkeCrsConstructor = {
-  deserialize(buffer: Uint8Array): ZkCompactPkeCrs;
-  safe_deserialize(buffer: Uint8Array, serialized_size_limit: bigint): ZkCompactPkeCrs;
-  from_config(config: unknown, max_num_bits: number): ZkCompactPkeCrs;
-  deserialize_from_public_params(buffer: Uint8Array): ZkCompactPkeCrs;
-  safe_deserialize_from_public_params(buffer: Uint8Array, serialized_size_limit: bigint): ZkCompactPkeCrs;
-};
-
 export type ZkCiphertextListBuilder = {
   push_boolean(data: boolean): void;
   push_u8(data: number): void;
@@ -98,7 +90,6 @@ export const MAX_UINT64: bigint = 18446744073709551615n; // 2^64 - 1
 export const MAX_UINT128: bigint = 340282366920938463463374607431768211455n; // 2^128 - 1
 export const MAX_UINT256: bigint = 115792089237316195423570985008687907853269984665640564039457584007913129640319n; // 2^256 - 1
 export const MAX_UINT160: bigint = 1461501637330902918203684832716283019655932542975n; // 2^160 - 1
-export const MAX_ENCRYPTABLE_BITS: number = 2048;
 
 // ===== CORE FUNCTIONS =====
 
@@ -174,14 +165,14 @@ export const zkPack = (items: EncryptableItem[], builder: ZkCiphertextListBuilde
     }
   }
 
-  if (totalBits > MAX_ENCRYPTABLE_BITS) {
+  if (totalBits > TFHE_RS_ZK_MAX_BITS) {
     throw new CofheError({
       code: CofheErrorCode.ZkPackFailed,
-      message: `Total bits ${totalBits} exceeds ${MAX_ENCRYPTABLE_BITS}`,
-      hint: `Ensure that the total bits of the items to encrypt does not exceed ${MAX_ENCRYPTABLE_BITS}`,
+      message: `Total bits ${totalBits} exceeds ${TFHE_RS_ZK_MAX_BITS}`,
+      hint: `Ensure that the total bits of the items to encrypt does not exceed ${TFHE_RS_ZK_MAX_BITS}`,
       context: {
         totalBits,
-        maxBits: MAX_ENCRYPTABLE_BITS,
+        maxBits: TFHE_RS_ZK_MAX_BITS,
         items,
       },
     });
@@ -231,7 +222,7 @@ export const zkProve = async (
         1 // ZkComputeLoad.Verify
       );
 
-      resolve(compactList.serialize());
+      resolve(compactList.safe_serialize(TFHE_RS_SAFE_SERIALIZATION_SIZE_LIMIT));
     }, 0);
   });
 };

package/core/fetchKeys.ts

@@ -1,5 +1,3 @@
-import { hardhat } from '@/chains';
-
 import { type CofheConfig, getCoFheUrlOrThrow } from './config.js';
 import { type KeysStorage } from './keyStore.js';
 

package/core/index.ts

@@ -43,6 +43,9 @@ export type {
   FheTypeValue,
   // Decryption types
   UnsealedItem,
+  DecryptPollCallbackFunction,
+  DecryptPollCallbackContext,
+  DecryptEndpoint,
   // Util types
   EncryptStepCallbackFunction as EncryptSetStateFn,
   EncryptStepCallbackContext,
@@ -85,7 +88,7 @@ export type {
 } from './encrypt/zkPackProveVerify.js';
 export { zkProveWithWorker } from './encrypt/zkPackProveVerify.js';
 
-// Contract addresses
+// Consts (contract addresses, serialized size limits)
 export {
   TASK_MANAGER_ADDRESS,
   MOCKS_ZK_VERIFIER_ADDRESS,
@@ -94,7 +97,12 @@ export {
   MOCKS_DECRYPT_RESULT_SIGNER_PRIVATE_KEY,
   MOCKS_THRESHOLD_NETWORK_ADDRESS,
   TEST_BED_ADDRESS,
+  TFHE_RS_ZK_MAX_BITS,
+  TFHE_RS_SAFE_SERIALIZATION_SIZE_LIMIT,
 } from './consts.js';
 
+// Decrypt result verification
+export { verifyDecryptResult } from './decrypt/verifyDecryptResult.js';
+
 // Utils
 export { fheTypeToString } from './utils.js';

package/core/keyStore.ts

@@ -2,6 +2,7 @@ import { createStore, type StoreApi } from 'zustand/vanilla';
 import { persist, createJSONStorage } from 'zustand/middleware';
 import { produce } from 'immer';
 import { type IStorage } from './types.js';
+import { TFHE_RS_KEY_VERSION } from './consts.js';
 
 // Type definitions
 type ChainRecord<T> = Record<string, T>;
@@ -23,6 +24,8 @@ export type KeysStorage = {
   rehydrateKeysStore: () => Promise<void>;
 };
 
+const KEYSTORE_NAME = `cofhesdk-keys-v${TFHE_RS_KEY_VERSION}`;
+
 function isValidPersistedState(state: unknown): state is KeysStore {
   if (state && typeof state === 'object') {
     if ('fhe' in state && 'crs' in state) {
@@ -94,7 +97,7 @@ export function createKeysStore(storage: IStorage | null): KeysStorage {
 
   const clearKeysStorage = async () => {
     if (storage) {
-      await storage.removeItem('cofhesdk-keys');
+      await storage.removeItem(KEYSTORE_NAME);
     }
     // If no storage, this is a no-op
   };
@@ -125,7 +128,7 @@ function createStoreWithPersit(storage: IStorage) {
       onRehydrateStorage: () => (_state?, _error?) => {
         if (_error) throw new Error(`onRehydrateStorage: Error rehydrating keys store: ${_error}`);
       },
-      name: 'cofhesdk-keys',
+      name: KEYSTORE_NAME,
       storage: createJSONStorage(() => storage),
       merge: (persistedState, currentState) => {
         const persisted = isValidPersistedState(persistedState) ? persistedState : DEFAULT_KEYS_STORE;