@cofhe/sdk 0.1.1 → 0.2.1

package/core/client.ts

@@ -1,59 +1,61 @@
-/* eslint-disable no-unused-vars */
 import type { CreateSelfPermitOptions, CreateSharingPermitOptions, ImportSharedPermitOptions } from '@/permits';
 
 import { createStore } from 'zustand/vanilla';
 import { type PublicClient, type WalletClient } from 'viem';
-import { fetchMultichainKeys } from './fetchKeys.js';
 import { CofhesdkError, CofhesdkErrorCode } from './error.js';
 import { EncryptInputsBuilder } from './encrypt/encryptInputsBuilder.js';
-import { type Result, ResultOk, resultWrapper } from './result.js';
 import { createKeysStore } from './keyStore.js';
 import { permits } from './permits.js';
 import { DecryptHandlesBuilder } from './decrypt/decryptHandleBuilder.js';
-import {
-  type CofhesdkClient,
-  type CofhesdkClientParams,
-  type CofhesdkClientConnectionState,
-  type EncryptableItem,
-  type FheTypes,
-  type CofhesdkClientPermits,
-} from './types.js';
 import { getPublicClientChainID, getWalletClientAccount } from './utils.js';
+import type {
+  CofhesdkClientConnectionState,
+  CofhesdkClientParams,
+  CofhesdkClient,
+  CofhesdkClientPermits,
+} from './clientTypes.js';
+import type { EncryptableItem, FheTypes } from './types.js';
+import type { CofhesdkConfig } from './config.js';
+
+export const InitialConnectStore: CofhesdkClientConnectionState = {
+  connected: false,
+  connecting: false,
+  connectError: undefined,
+  chainId: undefined,
+  account: undefined,
+  publicClient: undefined,
+  walletClient: undefined,
+};
 
 /**
  * Creates a CoFHE SDK client instance (base implementation)
  * @param {CofhesdkClientParams} opts - Initialization options including config and platform-specific serializers
  * @returns {CofhesdkClient} - The CoFHE SDK client instance
  */
-export function createCofhesdkClientBase(opts: CofhesdkClientParams): CofhesdkClient {
+export function createCofhesdkClientBase<TConfig extends CofhesdkConfig>(
+  opts: CofhesdkClientParams<TConfig>
+): CofhesdkClient<TConfig> {
   // Create keysStorage instance using configured storage
   const keysStorage = createKeysStore(opts.config.fheKeyStorage);
 
-  // refs captured in closure
-  let _publicClient: PublicClient | undefined = undefined;
-  let _walletClient: WalletClient | undefined = undefined;
-
   // Zustand store for reactive state management
-  const connectStore = createStore<CofhesdkClientConnectionState>(() => ({
-    connected: false,
-    connecting: false,
-    connectError: undefined,
-    chainId: undefined,
-    account: undefined,
-  }));
+
+  const connectStore = createStore<CofhesdkClientConnectionState>(() => InitialConnectStore);
+
+  // Minimal cancellation mechanism: incremented on each connect/disconnect.
+  // If a connect finishes after a disconnect, it must not overwrite the disconnected state.
+  let connectAttemptId = 0;
 
   // Helper to update state
   const updateConnectState = (partial: Partial<CofhesdkClientConnectionState>) => {
     connectStore.setState((state) => ({ ...state, ...partial }));
   };
 
-  // single-flight + abortable warmup
-  let _connectPromise: Promise<Result<boolean>> | undefined = undefined;
-
   // Called before any operation, throws of connection not yet established
   const _requireConnected = () => {
     const state = connectStore.getState();
-    const notConnected = !state.connected || !_publicClient || !_walletClient || !state.account || !state.chainId;
+    const notConnected =
+      !state.connected || !state.account || !state.chainId || !state.publicClient || !state.walletClient;
     if (notConnected) {
       throw new CofhesdkError({
         code: CofhesdkErrorCode.NotConnected,
@@ -63,74 +65,62 @@ export function createCofhesdkClientBase(opts: CofhesdkClientParams): CofhesdkCl
           connected: state.connected,
           account: state.account,
           chainId: state.chainId,
-          publicClient: _publicClient,
-          walletClient: _walletClient,
+          publicClient: state.publicClient,
+          walletClient: state.walletClient,
         },
       });
     }
   };
 
-  // INITIALIZATION
-
-  const keyFetchResult = resultWrapper(async () => {
-    // If configured, fetch keys for all supported chains
-    if (opts.config.fheKeysPrefetching === 'SUPPORTED_CHAINS') {
-      await fetchMultichainKeys(
-        opts.config,
-        0,
-        opts.tfhePublicKeyDeserializer,
-        opts.compactPkeCrsDeserializer,
-        keysStorage
-      );
-      return true;
-    }
-
-    return false;
-  });
-
   // LIFECYCLE
 
   async function connect(publicClient: PublicClient, walletClient: WalletClient) {
     const state = connectStore.getState();
 
     // Exit if already connected and clients are the same
-    if (state.connected && _publicClient === publicClient && _walletClient === walletClient) {
-      return Promise.resolve(ResultOk(true));
-    }
+    if (state.connected && state.publicClient === publicClient && state.walletClient === walletClient) return;
 
-    // Exit if already connecting
-    if (_connectPromise && _publicClient === publicClient && _walletClient === walletClient) {
-      return _connectPromise;
-    }
+    connectAttemptId += 1;
+    const localAttemptId = connectAttemptId;
 
     // Set connecting state
-    updateConnectState({ connecting: true, connectError: null, connected: false });
-
-    _connectPromise = resultWrapper(
-      // try
-      async () => {
-        _publicClient = publicClient;
-        _walletClient = walletClient;
-
-        const chainId = await getPublicClientChainID(publicClient);
-        const account = await getWalletClientAccount(walletClient);
-
-        updateConnectState({ connecting: false, connected: true, chainId, account });
-
-        return true;
-      },
-      // catch
-      (e) => {
-        updateConnectState({ connecting: false, connected: false, connectError: e });
-        return false;
-      },
-      // finally
-      () => {
-        _connectPromise = undefined;
-      }
-    );
-
-    return _connectPromise;
+    updateConnectState({
+      ...InitialConnectStore,
+      connecting: true,
+    });
+
+    // Fetch chainId and account
+    try {
+      const chainId = await getPublicClientChainID(publicClient);
+      const account = await getWalletClientAccount(walletClient);
+
+      // If a disconnect (or a newer connect) happened while awaiting, ignore this completion.
+      if (localAttemptId !== connectAttemptId) return;
+
+      updateConnectState({
+        connected: true,
+        connecting: false,
+        connectError: undefined,
+        chainId,
+        account,
+        publicClient,
+        walletClient,
+      });
+    } catch (e) {
+      // Ignore stale errors too.
+      if (localAttemptId !== connectAttemptId) return;
+
+      updateConnectState({
+        ...InitialConnectStore,
+        connectError: e,
+      });
+      throw e;
+    }
+  }
+
+  function disconnect() {
+    connectAttemptId += 1;
+    updateConnectState({ ...InitialConnectStore });
   }
 
   // CLIENT OPERATIONS
@@ -144,14 +134,15 @@ export function createCofhesdkClientBase(opts: CofhesdkClientParams): CofhesdkCl
       chainId: state.chainId ?? undefined,
 
       config: opts.config,
-      publicClient: _publicClient ?? undefined,
-      walletClient: _walletClient ?? undefined,
+      publicClient: state.publicClient ?? undefined,
+      walletClient: state.walletClient ?? undefined,
       zkvWalletClient: opts.config._internal?.zkvWalletClient,
 
       tfhePublicKeyDeserializer: opts.tfhePublicKeyDeserializer,
       compactPkeCrsDeserializer: opts.compactPkeCrsDeserializer,
       zkBuilderAndCrsGenerator: opts.zkBuilderAndCrsGenerator,
       initTfhe: opts.initTfhe,
+      zkProveWorkerFn: opts.zkProveWorkerFn,
 
       keysStorage,
 
@@ -169,8 +160,8 @@ export function createCofhesdkClientBase(opts: CofhesdkClientParams): CofhesdkCl
       account: state.account ?? undefined,
 
       config: opts.config,
-      publicClient: _publicClient ?? undefined,
-      walletClient: _walletClient ?? undefined,
+      publicClient: state.publicClient ?? undefined,
+      walletClient: state.walletClient ?? undefined,
 
       requireConnected: _requireConnected,
     });
@@ -204,67 +195,84 @@ export function createCofhesdkClientBase(opts: CofhesdkClientParams): CofhesdkCl
     subscribe: permits.subscribe,
 
     // Creation methods (require connection)
-    createSelf: async (options: CreateSelfPermitOptions) =>
-      resultWrapper(async () => {
-        _requireConnected();
-        return permits.createSelf(options, _publicClient!, _walletClient!);
-      }),
-
-    createSharing: async (options: CreateSharingPermitOptions) =>
-      resultWrapper(async () => {
-        _requireConnected();
-        return permits.createSharing(options, _publicClient!, _walletClient!);
-      }),
-
-    importShared: async (options: ImportSharedPermitOptions | any | string) =>
-      resultWrapper(async () => {
-        _requireConnected();
-        return permits.importShared(options, _publicClient!, _walletClient!);
-      }),
+    createSelf: async (
+      options: CreateSelfPermitOptions,
+      clients?: { publicClient: PublicClient; walletClient: WalletClient }
+    ) => {
+      _requireConnected();
+      const { publicClient, walletClient } = clients ?? connectStore.getState();
+      return permits.createSelf(options, publicClient!, walletClient!);
+    },
+
+    createSharing: async (
+      options: CreateSharingPermitOptions,
+      clients?: { publicClient: PublicClient; walletClient: WalletClient }
+    ) => {
+      _requireConnected();
+      const { publicClient, walletClient } = clients ?? connectStore.getState();
+      return permits.createSharing(options, publicClient!, walletClient!);
+    },
+
+    importShared: async (
+      options: ImportSharedPermitOptions | string,
+      clients?: { publicClient: PublicClient; walletClient: WalletClient }
+    ) => {
+      _requireConnected();
+      const { publicClient, walletClient } = clients ?? connectStore.getState();
+      return permits.importShared(options, publicClient!, walletClient!);
+    },
+
+    // Get or create methods (require connection)
+    getOrCreateSelfPermit: async (chainId?: number, account?: string, options?: CreateSelfPermitOptions) => {
+      _requireConnected();
+      const { chainId: _chainId, account: _account } = _getChainIdAndAccount(chainId, account);
+      const { publicClient, walletClient } = connectStore.getState();
+      return permits.getOrCreateSelfPermit(publicClient!, walletClient!, _chainId, _account, options);
+    },
+
+    getOrCreateSharingPermit: async (options: CreateSharingPermitOptions, chainId?: number, account?: string) => {
+      _requireConnected();
+      const { chainId: _chainId, account: _account } = _getChainIdAndAccount(chainId, account);
+      const { publicClient, walletClient } = connectStore.getState();
+      return permits.getOrCreateSharingPermit(publicClient!, walletClient!, options, _chainId, _account);
+    },
 
     // Retrieval methods (auto-fill chainId/account)
-    getPermit: async (hash: string, chainId?: number, account?: string) =>
-      resultWrapper(async () => {
-        const { chainId: _chainId, account: _account } = _getChainIdAndAccount(chainId, account);
-        return permits.getPermit(_chainId, _account, hash);
-      }),
-
-    getPermits: async (chainId?: number, account?: string) =>
-      resultWrapper(async () => {
-        const { chainId: _chainId, account: _account } = _getChainIdAndAccount(chainId, account);
-        return permits.getPermits(_chainId, _account);
-      }),
-
-    getActivePermit: async (chainId?: number, account?: string) =>
-      resultWrapper(async () => {
-        const { chainId: _chainId, account: _account } = _getChainIdAndAccount(chainId, account);
-        return permits.getActivePermit(_chainId, _account);
-      }),
-
-    getActivePermitHash: async (chainId?: number, account?: string) =>
-      resultWrapper(async () => {
-        const { chainId: _chainId, account: _account } = _getChainIdAndAccount(chainId, account);
-        return permits.getActivePermitHash(_chainId, _account);
-      }),
+    getPermit: async (hash: string, chainId?: number, account?: string) => {
+      const { chainId: _chainId, account: _account } = _getChainIdAndAccount(chainId, account);
+      return permits.getPermit(_chainId, _account, hash);
+    },
+
+    getPermits: async (chainId?: number, account?: string) => {
+      const { chainId: _chainId, account: _account } = _getChainIdAndAccount(chainId, account);
+      return permits.getPermits(_chainId, _account);
+    },
+
+    getActivePermit: async (chainId?: number, account?: string) => {
+      const { chainId: _chainId, account: _account } = _getChainIdAndAccount(chainId, account);
+      return permits.getActivePermit(_chainId, _account);
+    },
+
+    getActivePermitHash: async (chainId?: number, account?: string) => {
+      const { chainId: _chainId, account: _account } = _getChainIdAndAccount(chainId, account);
+      return permits.getActivePermitHash(_chainId, _account);
+    },
 
     // Mutation methods (auto-fill chainId/account)
-    selectActivePermit: async (hash: string, chainId?: number, account?: string) =>
-      resultWrapper(async () => {
-        const { chainId: _chainId, account: _account } = _getChainIdAndAccount(chainId, account);
-        return permits.selectActivePermit(_chainId, _account, hash);
-      }),
-
-    removePermit: async (hash: string, chainId?: number, account?: string) =>
-      resultWrapper(async () => {
-        const { chainId: _chainId, account: _account } = _getChainIdAndAccount(chainId, account);
-        return permits.removePermit(_chainId, _account, hash);
-      }),
-
-    removeActivePermit: async (chainId?: number, account?: string) =>
-      resultWrapper(async () => {
-        const { chainId: _chainId, account: _account } = _getChainIdAndAccount(chainId, account);
-        return permits.removeActivePermit(_chainId, _account);
-      }),
+    selectActivePermit: (hash: string, chainId?: number, account?: string) => {
+      const { chainId: _chainId, account: _account } = _getChainIdAndAccount(chainId, account);
+      return permits.selectActivePermit(_chainId, _account, hash);
+    },
+
+    removePermit: async (hash: string, chainId?: number, account?: string) => {
+      const { chainId: _chainId, account: _account } = _getChainIdAndAccount(chainId, account);
+      return permits.removePermit(_chainId, _account, hash);
+    },
+
+    removeActivePermit: async (chainId?: number, account?: string) => {
+      const { chainId: _chainId, account: _account } = _getChainIdAndAccount(chainId, account);
+      return permits.removeActivePermit(_chainId, _account);
+    },
 
     // Utils (no context needed)
     getHash: permits.getHash,
@@ -277,11 +285,6 @@ export function createCofhesdkClientBase(opts: CofhesdkClientParams): CofhesdkCl
     getSnapshot: connectStore.getState,
     subscribe: connectStore.subscribe,
 
-    // initialization results
-    initializationResults: {
-      keyFetchResult,
-    },
-
     // flags (read-only: reflect snapshot)
     get connected() {
       return connectStore.getState().connected;
@@ -294,6 +297,7 @@ export function createCofhesdkClientBase(opts: CofhesdkClientParams): CofhesdkCl
     config: opts.config,
 
     connect,
+    disconnect,
     encryptInputs,
     decryptHandle,
     permits: clientPermits,
@@ -302,7 +306,7 @@ export function createCofhesdkClientBase(opts: CofhesdkClientParams): CofhesdkCl
     // Example:
     // async encryptData(data: unknown) {
     //   requireConnected();
-    //   // Use _publicClient and _walletClient for implementation
+    //   // Use state.publicClient and state.walletClient for implementation
     // },
   };
 }

package/core/clientTypes.ts

@@ -0,0 +1,114 @@
+// TODO: Extract client types to its own file, keep this one as primitives
+import { type PublicClient, type WalletClient } from 'viem';
+import { type CofhesdkConfig, type CofhesdkEnvironment } from './config.js';
+import { type DecryptHandlesBuilder } from './decrypt/decryptHandleBuilder.js';
+import { type EncryptInputsBuilder } from './encrypt/encryptInputsBuilder.js';
+import { type ZkBuilderAndCrsGenerator, type ZkProveWorkerFunction } from './encrypt/zkPackProveVerify.js';
+import { type FheKeyDeserializer } from './fetchKeys.js';
+import { permits } from './permits.js';
+import type { EncryptableItem, FheTypes, TfheInitializer } from './types.js';
+import type { PermitUtils } from 'permits/permit.js';
+import type {
+  CreateSelfPermitOptions,
+  Permit,
+  CreateSharingPermitOptions,
+  ImportSharedPermitOptions,
+  SharingPermit,
+  RecipientPermit,
+  SelfPermit,
+} from 'permits/types.js';
+
+// CLIENT
+
+export type CofhesdkClient<TConfig extends CofhesdkConfig = CofhesdkConfig> = {
+  // --- state access ---
+  getSnapshot(): CofhesdkClientConnectionState;
+  subscribe(listener: Listener): () => void;
+
+  // --- convenience flags (read-only) ---
+  readonly connected: boolean;
+  readonly connecting: boolean;
+
+  // --- config & platform-specific ---
+  readonly config: TConfig;
+
+  connect(publicClient: PublicClient, walletClient: WalletClient): Promise<void>;
+  /**
+   * Clears the current connection state (account/chainId/clients) and marks the client as disconnected.
+   *
+   * This does not delete persisted permits or stored FHE keys; it only resets the in-memory connection.
+   */
+  disconnect(): void;
+  /**
+   * Types docstring
+   */
+  encryptInputs<T extends EncryptableItem[]>(inputs: [...T]): EncryptInputsBuilder<[...T]>;
+  decryptHandle<U extends FheTypes>(ctHash: bigint, utype: U): DecryptHandlesBuilder<U>;
+  permits: CofhesdkClientPermits;
+};
+
+export type CofhesdkClientConnectionState = {
+  connected: boolean;
+  connecting: boolean;
+  connectError: unknown | undefined;
+  chainId: number | undefined;
+  account: `0x${string}` | undefined;
+  publicClient: PublicClient | undefined;
+  walletClient: WalletClient | undefined;
+};
+
+type Listener = (snapshot: CofhesdkClientConnectionState) => void;
+
+export type CofhesdkClientPermitsClients = {
+  publicClient: PublicClient;
+  walletClient: WalletClient;
+};
+
+export type CofhesdkClientPermits = {
+  getSnapshot: typeof permits.getSnapshot;
+  subscribe: typeof permits.subscribe;
+
+  // Creation methods (require connection, no params)
+  createSelf: (options: CreateSelfPermitOptions, clients?: CofhesdkClientPermitsClients) => Promise<SelfPermit>;
+  createSharing: (
+    options: CreateSharingPermitOptions,
+    clients?: CofhesdkClientPermitsClients
+  ) => Promise<SharingPermit>;
+  importShared: (
+    options: ImportSharedPermitOptions | string,
+    clients?: CofhesdkClientPermitsClients
+  ) => Promise<RecipientPermit>;
+
+  // Retrieval methods (chainId/account optional)
+  getPermit: (hash: string, chainId?: number, account?: string) => Promise<Permit | undefined>;
+  getPermits: (chainId?: number, account?: string) => Promise<Record<string, Permit>>;
+  getActivePermit: (chainId?: number, account?: string) => Promise<Permit | undefined>;
+  getActivePermitHash: (chainId?: number, account?: string) => Promise<string | undefined>;
+
+  // Get or create methods (get active or create new, chainId/account optional)
+  getOrCreateSelfPermit: (chainId?: number, account?: string, options?: CreateSelfPermitOptions) => Promise<Permit>;
+  getOrCreateSharingPermit: (
+    options: CreateSharingPermitOptions,
+    chainId?: number,
+    account?: string
+  ) => Promise<Permit>;
+
+  // Mutation methods (chainId/account optional)
+  selectActivePermit: (hash: string, chainId?: number, account?: string) => void;
+  removePermit: (hash: string, chainId?: number, account?: string) => void;
+  removeActivePermit: (chainId?: number, account?: string) => void;
+
+  // Utils
+  getHash: typeof PermitUtils.getHash;
+  serialize: typeof PermitUtils.serialize;
+  deserialize: typeof PermitUtils.deserialize;
+};
+
+export type CofhesdkClientParams<TConfig extends CofhesdkConfig> = {
+  config: TConfig;
+  zkBuilderAndCrsGenerator: ZkBuilderAndCrsGenerator;
+  tfhePublicKeyDeserializer: FheKeyDeserializer;
+  compactPkeCrsDeserializer: FheKeyDeserializer;
+  initTfhe: TfheInitializer;
+  zkProveWorkerFn?: ZkProveWorkerFunction;
+};

package/core/config.test.ts

@@ -1,4 +1,3 @@
-/* eslint-disable no-unused-vars */
 import { sepolia, hardhat } from '@/chains';
 
 import { describe, it, expect, vi } from 'vitest';
@@ -52,6 +51,17 @@ describe('createCofhesdkConfigBase', () => {
     expect(getNestedValue(result, path)).toEqual(expectedValue);
   };
 
+  it('environment', () => {
+    expectInvalidConfigItem('environment', 'not-a-valid-environment');
+    expectInvalidConfigItem('environment', 123);
+    expectInvalidConfigItem('environment', {});
+
+    expectValidConfigItem('environment', 'node', 'node');
+    expectValidConfigItem('environment', 'hardhat', 'hardhat');
+    expectValidConfigItem('environment', 'web', 'web');
+    expectValidConfigItem('environment', 'react', 'react');
+  });
+
   it('supportedChains', () => {
     expectInvalidConfigItem('supportedChains', {});
     expectInvalidConfigItem('supportedChains', 'not-an-array');
@@ -62,16 +72,6 @@ describe('createCofhesdkConfigBase', () => {
     expectValidConfigItem('supportedChains', [sepolia, hardhat], [sepolia, hardhat]);
   });
 
-  it('fheKeysPrefetching', () => {
-    expectInvalidConfigItem('fheKeysPrefetching', 'invalid-option');
-    expectInvalidConfigItem('fheKeysPrefetching', 5);
-
-    expectValidConfigItem('fheKeysPrefetching', 'CONNECTED_CHAIN', 'CONNECTED_CHAIN');
-    expectValidConfigItem('fheKeysPrefetching', 'SUPPORTED_CHAINS', 'SUPPORTED_CHAINS');
-    expectValidConfigItem('fheKeysPrefetching', 'OFF', 'OFF');
-    expectValidConfigItem('fheKeysPrefetching', undefined, 'OFF');
-  });
-
   it('permitGeneration', () => {
     expectInvalidConfigItem('permitGeneration', 'not-a-boolean');
     expectInvalidConfigItem('permitGeneration', null);
@@ -126,6 +126,14 @@ describe('createCofhesdkConfigBase', () => {
     expect(getItemCalled).toBe(true);
     expect(setItemCalled).toBe(true);
     expect(removeItemCalled).toBe(true);
+
+    const invalidStorageNotAFunction = {
+      getItem: 'not-a-function',
+      setItem: 'not-a-function',
+      removeItem: 'not-a-function',
+    };
+
+    expectInvalidConfigItem('fheKeyStorage', invalidStorageNotAFunction);
   });
 
   it('mocks', () => {
@@ -144,6 +152,17 @@ describe('createCofhesdkConfigBase', () => {
     expectValidConfigItem('mocks.sealOutputDelay', 1000, 1000);
   });
 
+  it('useWorkers', () => {
+    expectInvalidConfigItem('useWorkers', 'not-a-boolean');
+    expectInvalidConfigItem('useWorkers', null);
+    expectInvalidConfigItem('useWorkers', 123);
+    expectInvalidConfigItem('useWorkers', {});
+
+    expectValidConfigItem('useWorkers', true, true);
+    expectValidConfigItem('useWorkers', false, false);
+    expectValidConfigItem('useWorkers', undefined, true); // defaults to true
+  });
+
   it('should get config item', () => {
     const config: CofhesdkInputConfig = {
       supportedChains: [sepolia],