@cofhe/sdk 0.7.1 → 0.8.0

package/chains/test/chains.test.ts

@@ -1,12 +1,24 @@
 import { describe, it, expect } from 'vitest';
-import { sepolia, arbSepolia, baseSepolia, hardhat, chains, getChainById, getChainByName } from '../index.js';
+import {
+  sepolia,
+  arbSepolia,
+  baseSepolia,
+  ethereum,
+  arbitrum,
+  hardhat,
+  chains,
+  getChainById,
+  getChainByName,
+} from '../index.js';
 
 describe('Chains', () => {
   it('should export all chains', () => {
-    expect(Object.keys(chains)).toHaveLength(6);
+    expect(Object.keys(chains)).toHaveLength(8);
     expect(chains).toHaveProperty('sepolia');
     expect(chains).toHaveProperty('arbSepolia');
     expect(chains).toHaveProperty('baseSepolia');
+    expect(chains).toHaveProperty('ethereum');
+    expect(chains).toHaveProperty('arbitrum');
     expect(chains).toHaveProperty('hardhat');
     expect(chains).toHaveProperty('localcofhe');
     expect(chains).toHaveProperty('stagingCofhe');
@@ -17,6 +29,11 @@ describe('Chains', () => {
     expect(sepolia.name).toBe('Sepolia');
     expect(sepolia.environment).toBe('TESTNET');
 
+    expect(ethereum.id).toBe(1);
+    expect(ethereum.environment).toBe('MAINNET');
+    expect(arbitrum.id).toBe(42161);
+    expect(arbitrum.environment).toBe('MAINNET');
+
     expect(hardhat.id).toBe(31337);
     expect(hardhat.name).toBe('Hardhat');
     expect(hardhat.environment).toBe('MOCK');
@@ -24,6 +41,8 @@ describe('Chains', () => {
 
   it('should find chains by ID', () => {
     expect(getChainById(11155111)).toBe(sepolia);
+    expect(getChainById(1)).toBe(ethereum);
+    expect(getChainById(42161)).toBe(arbitrum);
     expect(getChainById(31337)).toBe(hardhat);
     expect(getChainById(999999)).toBeUndefined();
   });
@@ -31,12 +50,14 @@ describe('Chains', () => {
   it('should find chains by name', () => {
     expect(getChainByName('sepolia')).toBe(sepolia);
     expect(getChainByName('Sepolia')).toBe(sepolia);
+    expect(getChainByName('ethereum')).toBe(ethereum);
+    expect(getChainByName('arbitrum')).toBe(arbitrum);
     expect(getChainByName('hardhat')).toBe(hardhat);
     expect(getChainByName('nonexistent')).toBeUndefined();
   });
 
   it('should validate chain properties', () => {
-    const allChains = [sepolia, arbSepolia, baseSepolia, hardhat];
+    const allChains = [sepolia, arbSepolia, baseSepolia, ethereum, arbitrum, hardhat];
 
     allChains.forEach((chain) => {
       expect(typeof chain.id).toBe('number');

package/core/acps.ts

@@ -11,26 +11,17 @@ import {
   type IncomingShare,
   type SharingACP,
   type ACPHashFields,
+  type ACPAccessStatus,
+  ACP_REVOKER_ABI,
+  ACP_SHARE_REGISTRY_ABI,
+  toChainShare,
+  computeShareId,
+  shareIdOfChainShare,
+  getAclServedAddresses,
+  clearAclServedAddresses,
 } from '@/acps';
 
-import {
-  type Hex,
-  type PublicClient,
-  type WalletClient,
-  encodeAbiParameters,
-  keccak256,
-  parseAbi,
-  zeroAddress,
-} from 'viem';
-
-import { TASK_MANAGER_ADDRESS } from './consts.js';
-
-// ACP default revoker (timestamp-based revocation) — interface shared by all revokers
-const ACP_VALIDATOR_ABI = parseAbi([
-  'function revokeSingle(uint256 id)',
-  'function revokeAllExisting()',
-  'function disabled(address issuer, uint256 id) view returns (bool)',
-]);
+import { type Hex, type PublicClient, type WalletClient, zeroAddress } from 'viem';
 
 // HELPERS
 
@@ -51,9 +42,10 @@ const storeActiveACP = async (acp: ACP, publicClient: any, walletClient: any) =>
 };
 
 // Generic function to handle acp creation with error handling.
-// `activate` controls whether the new acp becomes the issuer's active acp — true for
-// self/imported acps (the connected user decrypts with them), false for sharing acps (those
-// are delegated to a recipient and are never the issuer's own active acp).
+// `activate` controls whether the new acp becomes the connected account's active acp — true for
+// self acps and, by default, imported acps (the connected user decrypts with them); false for
+// sharing acps (delegated to a recipient, never the issuer's own active acp) and for imports
+// made with `activate: false`.
 const createACPWithSign = async <T, TACP extends ACP>(
   options: T,
   publicClient: PublicClient,
@@ -96,12 +88,23 @@ const createSharing = async (
   return createACPWithSign(options, publicClient, walletClient, ACPUtils.createSharingAndSign, false);
 };
 
+/** Options for importing a shared acp. */
+export type ImportActivationOptions = {
+  /**
+   * Make the imported acp the active one. Defaults to `true`.
+   * With `false` the acp is only stored: decrypts that use the active acp keep using the current one,
+   * and the imported acp is used only where it is passed explicitly (e.g. `.withACP(acp)`).
+   */
+  activate?: boolean;
+};
+
 const importShared = async (
   options: ImportSharedACPOptions | string,
   publicClient: PublicClient,
-  walletClient: WalletClient
+  walletClient: WalletClient,
+  { activate = true }: ImportActivationOptions = {}
 ): Promise<RecipientACP> => {
-  return createACPWithSign(options, publicClient, walletClient, ACPUtils.importSharedAndSign);
+  return createACPWithSign(options, publicClient, walletClient, ACPUtils.importSharedAndSign, activate);
 };
 
 // ACP UTILS
@@ -263,65 +266,6 @@ const applyACPDefaults = <
   return result;
 };
 
-// ACL-SERVED ADDRESSES (defaultRevokerContract / shareRegistry)
-
-const ACL_SERVED_ADDRESSES_ABI = parseAbi([
-  'function acl() view returns (address)',
-  'function defaultRevokerContract() view returns (address)',
-  'function shareRegistry() view returns (address)',
-]);
-
-export interface AclServedAddresses {
-  defaultRevoker?: Hex;
-  shareRegistry?: Hex;
-}
-
-const aclServedAddressesCache = new Map<number, AclServedAddresses>();
-
-/** Test hook: forget resolved addresses (e.g. between redeployments on one chainId). */
-const clearAclServedAddresses = () => aclServedAddressesCache.clear();
-
-/**
- * The ACP infrastructure addresses the chain's ACL serves (TaskManager -> acl()
- * -> getters). Zero addresses and pre-upgrade ACLs (getters absent -> revert)
- * resolve to `undefined` — callers fall back to `acp.*` config.
- *
- * Resolutions are cached per chainId. A failure to reach the TaskManager (network
- * error, no CoFHE deployment) is NOT cached, so a transient outage does not pin
- * an empty result for the whole session.
- */
-const getAclServedAddresses = async (publicClient: PublicClient, chainId: number): Promise<AclServedAddresses> => {
-  const cached = aclServedAddressesCache.get(chainId);
-  if (cached != null) return cached;
-
-  let aclAddress: Hex;
-  try {
-    aclAddress = await publicClient.readContract({
-      address: TASK_MANAGER_ADDRESS,
-      abi: ACL_SERVED_ADDRESSES_ABI,
-      functionName: 'acl',
-    });
-  } catch {
-    return {};
-  }
-
-  const [defaultRevoker, shareRegistry] = await Promise.all([
-    publicClient
-      .readContract({ address: aclAddress, abi: ACL_SERVED_ADDRESSES_ABI, functionName: 'defaultRevokerContract' })
-      .catch(() => undefined),
-    publicClient
-      .readContract({ address: aclAddress, abi: ACL_SERVED_ADDRESSES_ABI, functionName: 'shareRegistry' })
-      .catch(() => undefined),
-  ]);
-
-  const resolved: AclServedAddresses = {
-    defaultRevoker: defaultRevoker != null && defaultRevoker !== zeroAddress ? defaultRevoker : undefined,
-    shareRegistry: shareRegistry != null && shareRegistry !== zeroAddress ? shareRegistry : undefined,
-  };
-  aclServedAddressesCache.set(chainId, resolved);
-  return resolved;
-};
-
 /**
  * `applyACPDefaults` with the ACL consulted for the default revoker when
  * `acp.defaultRevoker` config does not name one for this chain — explicit
@@ -374,7 +318,7 @@ const revokeACP = async (acp: ACP, walletClient: WalletClient): Promise<Hex> =>
 
   return walletClient.writeContract({
     address: acp.revokerContract,
-    abi: ACP_VALIDATOR_ABI,
+    abi: ACP_REVOKER_ABI,
     functionName: 'revokeSingle',
     args: [BigInt(acp.revokerData)],
     account: walletClient.account,
@@ -410,7 +354,7 @@ const revokeAllACPs = async (
 
   return walletClient.writeContract({
     address: revoker,
-    abi: ACP_VALIDATOR_ABI,
+    abi: ACP_REVOKER_ABI,
     functionName: 'revokeAllExisting',
     args: [],
     account: walletClient.account,
@@ -418,6 +362,20 @@ const revokeAllACPs = async (
   });
 };
 
+/**
+ * The ACP on-chain status as a value instead of a revert (see `ACPAccessStatus`): validity
+ * without a handle, and whether the ACP may read `handle` with one.
+ */
+const checkAccess = async (acp: ACP, publicClient: PublicClient, handle?: bigint | Hex): Promise<ACPAccessStatus> => {
+  // The issuer half of a share carries no recipient signature (the recipient adds it on import), so
+  // the ACL check would always fail it. What can change for the issuer is expiry and revocation.
+  if (acp.type === 'sharing' && handle === undefined) {
+    if (ACPUtils.isExpired(acp)) return 'expired';
+    return (await isACPRevoked(acp, publicClient)) ? 'revoked' : 'valid';
+  }
+  return ACPUtils.checkAccessOnChain(acp, publicClient, handle);
+};
+
 /**
  * Check whether an ACP has been revoked (or is otherwise disabled) by its
  * revoker. Returns false for acps without a revoker (not revocable).
@@ -426,65 +384,12 @@ const isACPRevoked = async (acp: ACP, publicClient: PublicClient): Promise<boole
   if (acp.revokerContract === zeroAddress || acp.revokerData === 0) return false;
   return publicClient.readContract({
     address: acp.revokerContract,
-    abi: ACP_VALIDATOR_ABI,
+    abi: ACP_REVOKER_ABI,
     functionName: 'disabled',
     args: [acp.issuer, BigInt(acp.revokerData)],
   });
 };
 
-// SHARE (on-chain, via the ACPShareRegistry)
-
-const ACP_SHARE_REGISTRY_ABI = parseAbi([
-  'struct ACP { address issuer; uint64 expiration; address recipient; uint256 revokerData; address revokerContract; uint8 scope; address[] contracts; bytes32[] handles; bytes32 sealingKey; bytes issuerSignature; bytes recipientSignature; }',
-  'function share(ACP calldata acp) external returns (bytes32)',
-  'function removeShare(bytes32 shareId) external',
-  'function sharesFor(address recipient) external view returns (ACP[] memory)',
-  'function getShare(bytes32 shareId) external view returns (ACP memory)',
-  'function isShareValid(bytes32 shareId) external view returns (bool)',
-]);
-
-const ACP_TUPLE = [
-  {
-    type: 'tuple',
-    components: [
-      { name: 'issuer', type: 'address' },
-      { name: 'expiration', type: 'uint64' },
-      { name: 'recipient', type: 'address' },
-      { name: 'revokerData', type: 'uint256' },
-      { name: 'revokerContract', type: 'address' },
-      { name: 'scope', type: 'uint8' },
-      { name: 'contracts', type: 'address[]' },
-      { name: 'handles', type: 'bytes32[]' },
-      { name: 'sealingKey', type: 'bytes32' },
-      { name: 'issuerSignature', type: 'bytes' },
-      { name: 'recipientSignature', type: 'bytes' },
-    ],
-  },
-] as const;
-
-const ZERO_BYTES32 = `0x${'0'.repeat(64)}` as Hex;
-
-/** The on-chain payload for a sharing ACP: recipient-side fields empty. */
-const toChainShare = (acp: ACP) => {
-  // Same public struct as the off-chain export flow (ACPUtils.getPublic), with
-  // the recipient-side fields blanked — the recipient supplies them at import —
-  // and uint fields widened for the ABI encoder.
-  const pub = ACPUtils.getPublic(acp, true);
-  return {
-    ...pub,
-    expiration: BigInt(pub.expiration),
-    revokerData: BigInt(pub.revokerData),
-    sealingKey: ZERO_BYTES32,
-    recipientSignature: '0x' as Hex,
-  };
-};
-
-/** Mirrors the registry's `keccak256(abi.encode(acp))` share id. */
-const computeShareId = (acp: ACP): Hex => {
-  const p = toChainShare(acp);
-  return keccak256(encodeAbiParameters(ACP_TUPLE, [p]));
-};
-
 /**
  * Post a signed sharing ACP to the on-chain share registry for its recipient
  * to discover and import — the on-chain alternative to `export()`.
@@ -531,7 +436,7 @@ const getIncomingShares = async (
   });
 
   return raw.map((s) => ({
-    shareId: keccak256(encodeAbiParameters(ACP_TUPLE, [s])),
+    shareId: shareIdOfChainShare(s),
     issuer: s.issuer,
     expiration: Number(s.expiration),
     recipient: s.recipient,
@@ -546,16 +451,18 @@ const getIncomingShares = async (
 
 /**
  * Import a share read from the registry: fills the recipient's sealing key,
- * signs, stores and activates — the on-chain counterpart of importing an
- * exported JSON blob. The share stays on-chain until dismissed.
+ * signs, stores and (unless `activate: false`) activates — the on-chain
+ * counterpart of importing an exported JSON blob. The share stays on-chain
+ * until dismissed.
  */
 const importFromChain = async (
   share: IncomingShare,
   publicClient: PublicClient,
-  walletClient: WalletClient
+  walletClient: WalletClient,
+  importOptions: ImportActivationOptions = {}
 ): Promise<RecipientACP> => {
   const { shareId: _shareId, ...options } = share;
-  return importShared({ ...options, type: 'sharing' }, publicClient, walletClient);
+  return importShared({ ...options, type: 'sharing' }, publicClient, walletClient, importOptions);
 };
 
 /** Remove a share from the registry (issuer retracts / recipient dismisses). */
@@ -608,6 +515,7 @@ export const acps = {
   revokeACP,
   revokeAllACPs,
   isACPRevoked,
+  checkAccess,
 
   shareOnChain,
   getIncomingShares,

package/core/client.ts

@@ -156,6 +156,7 @@ export function createCofheClientBase<TConfig extends CofheConfig>(
       zkBuilderAndCrsGenerator: opts.zkBuilderAndCrsGenerator,
       initTfhe: opts.initTfhe,
       zkProveWorkerFn: opts.zkProveWorkerFn,
+      beforeMainThreadProve: opts.beforeMainThreadProve,
 
       keysStorage,
 
@@ -259,13 +260,11 @@ export function createCofheClientBase<TConfig extends CofheConfig>(
       );
     },
 
-    importShared: async (
-      options: ImportSharedACPOptions | string,
-      clients?: { publicClient: PublicClient; walletClient: WalletClient }
-    ) => {
+    importShared: async (options: ImportSharedACPOptions | string, importOptions = {}) => {
       _requireConnected();
-      const { publicClient, walletClient } = clients ?? connectStore.getState();
-      return acps.importShared(options, publicClient!, walletClient!);
+      const { activate, ...clients } = importOptions;
+      const { publicClient, walletClient } = clients.publicClient ? clients : connectStore.getState();
+      return acps.importShared(options, publicClient!, walletClient!, { activate });
     },
 
     // Get or create methods (require connection)
@@ -314,6 +313,12 @@ export function createCofheClientBase<TConfig extends CofheConfig>(
       return acps.isACPRevoked(acp, publicClient!);
     },
 
+    checkAccess: async (acp, handle) => {
+      _requireConnected();
+      const { publicClient } = connectStore.getState();
+      return acps.checkAccess(acp, publicClient!, handle);
+    },
+
     // On-chain sharing (require connection + config.acp.sharingRegistry)
     shareOnChain: async (acp) => {
       _requireConnected();
@@ -332,10 +337,10 @@ export function createCofheClientBase<TConfig extends CofheConfig>(
       return acps.getIncomingShares(publicClient!, sharingRegistry, account);
     },
 
-    importFromChain: async (share) => {
+    importFromChain: async (share, options) => {
       _requireConnected();
       const { publicClient, walletClient } = connectStore.getState();
-      return acps.importFromChain(share, publicClient!, walletClient!);
+      return acps.importFromChain(share, publicClient!, walletClient!, options);
     },
 
     dismissShare: async (shareId) => {

package/core/clientTypes.ts

@@ -6,10 +6,11 @@ import { type DecryptForTxBuilderUnset } from './decrypt/decryptForTxBuilder.js'
 import { type EncryptInputsBuilderUnset } from './encrypt/encryptInputsBuilder.js';
 import { type ZkBuilderAndCrsGenerator, type ZkProveWorkerFunction } from './encrypt/zkPackProveVerify.js';
 import { type FheKeyDeserializer } from './fetchKeys.js';
-import { acps } from './acps.js';
+import { acps, type ImportActivationOptions } from './acps.js';
 import type { EncryptableItem, FheTypes, TfheInitializer } from './types.js';
 import type { ACPUtils } from 'acps/acp.js';
 import type {
+  ACPAccessStatus,
   CreateSelfACPOptions,
   ACP,
   CreateSharingACPOptions,
@@ -74,6 +75,13 @@ export type CofheClientACPsClients = {
   walletClient: WalletClient;
 };
 
+/**
+ * Options for `importShared`. `activate: false` stores the imported acp without making it the active
+ * one (default `true`). Clients default to the connected ones; pass both to override.
+ */
+export type CofheClientImportSharedOptions = ImportActivationOptions &
+  (CofheClientACPsClients | { publicClient?: undefined; walletClient?: undefined });
+
 export type CofheClientACPs = {
   getSnapshot: typeof acps.getSnapshot;
   subscribe: typeof acps.subscribe;
@@ -81,7 +89,10 @@ export type CofheClientACPs = {
   // Creation methods (require connection, no params)
   createSelf: (options: CreateSelfACPOptions, clients?: CofheClientACPsClients) => Promise<SelfACP>;
   createSharing: (options: CreateSharingACPOptions, clients?: CofheClientACPsClients) => Promise<SharingACP>;
-  importShared: (options: ImportSharedACPOptions | string, clients?: CofheClientACPsClients) => Promise<RecipientACP>;
+  importShared: (
+    options: ImportSharedACPOptions | string,
+    importOptions?: CofheClientImportSharedOptions
+  ) => Promise<RecipientACP>;
 
   // Retrieval methods (chainId/account optional)
   getACP: (hash: string, chainId?: number, account?: string) => ACP | undefined;
@@ -102,13 +113,20 @@ export type CofheClientACPs = {
   revokeACP: (acp: ACP) => Promise<`0x${string}`>;
   revokeAllACPs: (revokerContract?: `0x${string}`) => Promise<`0x${string}`>;
   isACPRevoked: (acp: ACP) => Promise<boolean>;
+  /**
+   * The ACP on-chain status as a value instead of a revert. Without a handle: 'valid', or
+   * 'expired' / 'revoked' / an invalid signature. With a handle, a valid ACP reports
+   * 'allowed', 'out-of-scope' or 'issuer-not-allowed'. For the issuer copy of a share (type
+   * 'sharing', no recipient signature yet) the handle-less check covers expiry and revocation only.
+   */
+  checkAccess: (acp: ACP, handle?: bigint | `0x${string}`) => Promise<ACPAccessStatus>;
 
   /** Post a signed sharing ACP to the on-chain share registry (ACL-served, or config `acp.sharingRegistry`). */
   shareOnChain: (acp: ACP) => Promise<{ txHash: `0x${string}`; shareId: `0x${string}` }>;
   /** Importable shares addressed to the connected account (unexpired, not revoked). */
   getIncomingShares: () => Promise<IncomingShare[]>;
-  /** Import a share read from the registry: sign as recipient, store and activate. */
-  importFromChain: (share: IncomingShare) => Promise<RecipientACP>;
+  /** Import a share read from the registry: sign as recipient, store and (unless `activate: false`) activate. */
+  importFromChain: (share: IncomingShare, options?: ImportActivationOptions) => Promise<RecipientACP>;
   /** Recipient-side: remove a share from the registry (after import, or to decline). */
   dismissShare: (shareId: `0x${string}`) => Promise<`0x${string}`>;
   /** Issuer-side: retract a pending share from the registry. */
@@ -128,4 +146,11 @@ export type CofheClientParams<TConfig extends CofheConfig> = {
   compactPkeCrsDeserializer: FheKeyDeserializer;
   initTfhe: TfheInitializer;
   zkProveWorkerFn?: ZkProveWorkerFunction;
+  /**
+   * Runs right before a ZK proof is generated on the calling thread, i.e. only
+   * when the worker path is disabled, unavailable or has failed. Platforms use
+   * it for setup that only pays off when that thread does the proving (web
+   * starts tfhe's rayon thread pool here).
+   */
+  beforeMainThreadProve?: () => Promise<void>;
 };

package/core/config.ts

@@ -3,7 +3,7 @@ import { type CofheChain } from '@/chains';
 import { z } from 'zod';
 import { type WalletClient } from 'viem';
 import { CofheError, CofheErrorCode } from './error.js';
-import { type IStorage } from './types.js';
+import { type IStorage, type TfheThreadsSetting } from './types.js';
 
 export type CofheEnvironment = 'node' | 'hardhat' | 'web' | 'react';
 
@@ -29,6 +29,25 @@ export type CofheConfig = {
    * Default: true
    */
   useWorkers: boolean;
+  /**
+   * How many threads tfhe's rayon pool may use for ZK proof generation
+   * (web platform only; ignored on node/hardhat).
+   *
+   * Proving dominates the cost of `encryptInputs`, and tfhe's wasm is built with
+   * rayon, so spreading it across threads is a large win — roughly 5x on an
+   * 8-thread pool.
+   *
+   * Requires a cross-origin-isolated page (`Cross-Origin-Opener-Policy: same-origin`
+   * plus `Cross-Origin-Embedder-Policy: require-corp` or `credentialless`). When the
+   * page isn't isolated the pool is skipped and proving stays single-threaded, so
+   * this is safe to leave on.
+   *
+   * `'auto'` derives the count from `navigator.hardwareConcurrency` (capped, since
+   * the speedup flattens out); a number requests exactly that many; `false` keeps
+   * tfhe single-threaded.
+   * Default: 'auto'
+   */
+  tfheThreads: TfheThreadsSetting;
   /** ACP acp defaults, applied when creating acps (user options always win) */
   acp: {
     /**
@@ -108,6 +127,11 @@ export const CofheConfigSchema = z.object({
     .default(null),
   /** Whether to use Web Workers for ZK proof generation (web platform only) */
   useWorkers: z.boolean().optional().default(true),
+  /** How many threads tfhe's rayon pool may use for ZK proof generation (web platform only) */
+  tfheThreads: z
+    .union([z.number().int().positive(), z.literal('auto'), z.literal(false)])
+    .optional()
+    .default('auto'),
   /** ACP acp defaults */
   acp: z
     .object({

package/core/consts.ts

@@ -1,5 +1,17 @@
-/** Main Task Manager contract address */
-export const TASK_MANAGER_ADDRESS = '0xeA30c4B8b44078Bbf8a6ef5b9f1eC1626C7848D9' as const;
+/**
+ * CoFHEAddressBook contract address - the same on every chain. Resolves `TASK_MANAGER_ID` to the
+ * chain's Task Manager. Keep in sync with `COFHE_ADDRESS_BOOK` in `@fhenixprotocol/cofhe-contracts/FHE.sol`.
+ */
+export const COFHE_ADDRESS_BOOK_ADDRESS = '0xC0F4e00E531a2B086492Ae3DCC1515038307196b' as const;
+
+/** Id of the Task Manager in the CoFHEAddressBook. Keep in sync with `TASK_MANAGER_ID` in FHE.sol */
+export const TASK_MANAGER_ID = 1n;
+
+/**
+ * Mock Task Manager contract address (used for testing). Registered in the mock CoFHEAddressBook;
+ * SDK reads still resolve the Task Manager through the book.
+ */
+export const MOCKS_TASK_MANAGER_ADDRESS = '0x0000000000000000000000000000000000005000' as const;
 
 /** Mock ZK Verifier contract address (used for testing) */
 export const MOCKS_ZK_VERIFIER_ADDRESS = '0x0000000000000000000000000000000000005001' as const;

package/core/decrypt/verifyDecryptResult.ts

@@ -9,7 +9,7 @@ import {
   type Hex,
   type PublicClient,
 } from 'viem';
-import { TASK_MANAGER_ADDRESS } from '../consts.js';
+import { getTaskManagerAddress } from '../taskManager.js';
 
 const decryptResultSignerAbi = parseAbi(['function decryptResultSigner() view returns (address)']);
 const UINT_TYPE_MASK = 0x7fn;
@@ -33,7 +33,8 @@ const buildDecryptResultHash = (ctHash: bigint, cleartext: bigint, chainId: numb
  * 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`).
+ * The only on-chain reads performed are the Task Manager lookup in the CoFHEAddressBook (cached
+ * per chain) and `TaskManager.decryptResultSigner()` (via `eth_call`).
  *
  * Works with both production and mock deployments.
  */
@@ -45,7 +46,7 @@ export async function verifyDecryptResult(
 ): Promise<boolean> {
   const chainId = publicClient.chain?.id ?? (await publicClient.getChainId());
   const expectedSigner = await publicClient.readContract({
-    address: TASK_MANAGER_ADDRESS,
+    address: await getTaskManagerAddress(publicClient, chainId),
     abi: decryptResultSignerAbi,
     functionName: 'decryptResultSigner',
     args: [],

package/core/encrypt/encryptInputsBuilder.ts

@@ -37,6 +37,7 @@ type EncryptInputsBuilderParams<T extends EncryptableItem[]> = BaseBuilderParams
   zkBuilderAndCrsGenerator: ZkBuilderAndCrsGenerator | undefined;
   initTfhe: TfheInitializer | undefined;
   zkProveWorkerFn: ZkProveWorkerFunction | undefined;
+  beforeMainThreadProve?: (() => Promise<void>) | undefined;
 
   keysStorage: KeysStorage | undefined;
 };
@@ -69,6 +70,7 @@ export class EncryptInputsBuilder<T extends EncryptableItem[]> extends BaseBuild
   private zkBuilderAndCrsGenerator: ZkBuilderAndCrsGenerator;
   private initTfhe: TfheInitializer | undefined;
   private zkProveWorkerFn: ZkProveWorkerFunction | undefined;
+  private beforeMainThreadProve: (() => Promise<void>) | undefined;
 
   private keysStorage: KeysStorage | undefined;
 
@@ -143,6 +145,9 @@ export class EncryptInputsBuilder<T extends EncryptableItem[]> extends BaseBuild
     // Optional zkProve worker function, will be used on web if useWorkers is true and worker function is provided
     this.zkProveWorkerFn = params.zkProveWorkerFn;
 
+    // Optional platform setup that only matters when this thread generates the proof
+    this.beforeMainThreadProve = params.beforeMainThreadProve;
+
     // Keys storage is used to store the FHE key and CRS
     this.keysStorage = params.keysStorage;
 
@@ -582,6 +587,7 @@ export class EncryptInputsBuilder<T extends EncryptableItem[]> extends BaseBuild
 
     if (proof == null) {
       // Use main thread directly (workers disabled or unavailable)
+      await this.beforeMainThreadProve?.();
       proof = await zkProve(zkBuilder, zkCrs, metadata);
       usedWorker = false;
     }

package/core/encrypt/zkPackProveVerify.ts

@@ -1,6 +1,6 @@
 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 { type EncryptableItem, FheTypes, type TfheThreadsSetting } from '../types';
 import { toBigIntOrThrow, validateBigIntInRange, toHexString, hexToBytes } from '../utils';
 
 // ===== TYPES =====
@@ -29,6 +29,12 @@ export interface ZkProveWorkerRequest {
     data: any;
   }>;
   metadata: number[]; // Uint8Array serialized as array
+  /**
+   * Forwarded from `config.tfheThreads`. The worker applies it the first time it
+   * initialises tfhe — `initThreadPool` can only run once per wasm instance, so
+   * later requests carrying a different value have no effect.
+   */
+  tfheThreads?: TfheThreadsSetting;
 }
 
 /**

package/core/error.ts

@@ -43,6 +43,7 @@ export enum CofheErrorCode {
   PublicWalletGetChainIdFailed = 'PUBLIC_WALLET_GET_CHAIN_ID_FAILED',
   PublicWalletGetAddressesFailed = 'PUBLIC_WALLET_GET_ADDRESSES_FAILED',
   RehydrateKeysStoreFailed = 'REHYDRATE_KEYS_STORE_FAILED',
+  TaskManagerUnresolved = 'TASK_MANAGER_UNRESOLVED',
 
   // Threshold-network stable error codes (decrypt/sealoutput), see API-RESPONSES.md
   BadRequest = 'BAD_REQUEST',

package/core/fetchKeys.ts

@@ -77,7 +77,7 @@ const fetchCrs = async (
   keysStorage?: KeysStorage | null
 ): Promise<[string, boolean]> => {
   // Escape if key already exists
-  const storedKey = keysStorage?.getCrs(chainId);
+  const storedKey = keysStorage?.getCrs(chainId, securityZone);
   const [storedKeyValid] = checkKeyValidity(storedKey, compactPkeCrsDeserializer);
   if (storedKeyValid) return [storedKey!, false];
 
@@ -108,7 +108,7 @@ const fetchCrs = async (
     throw new Error(`Error serializing CRS; ${err}`);
   }
 
-  keysStorage?.setCrs(chainId, crs_data);
+  keysStorage?.setCrs(chainId, securityZone, crs_data);
 
   return [crs_data, true];
 };