encrypt-worker.js

import initWasm, { encrypt_serialize_u64_radix_flat_wasm } from './concrete-ml-extensions-wasm/concrete_ml_extensions_wasm.js';

let clientKey;
let wasmInitialized = false;

self.onmessage = async function(e) {
    if (e.data.type === 'init') {
        try {
            await initWasm();
            wasmInitialized = true;
            clientKey = e.data.clientKey;
            console.log(`[Encrypt] Client key loaded: ${clientKey.length} bytes`);
            self.postMessage({ type: 'ready' });
        } catch (error) {
            console.error('[Encrypt] WASM initialization failed:', error.message);
            self.postMessage({ type: 'error', error: `WASM initialization failed: ${error.message}` });
        }
    } else if (e.data.type === 'encrypt') {
        if (!wasmInitialized) {
            console.error('[Encrypt] Attempted encryption before WASM initialization');
            self.postMessage({ type: 'error', error: 'WASM not initialized' });
            return;
        }
        try {
            const tokenIds = e.data.tokenIds;
            console.log(`[Encrypt] Starting encryption for ${tokenIds.length} token IDs`);
            
            // Convert token IDs to BigUint64Array for FHE encryption
            const bigArr = new BigUint64Array(tokenIds.map(id => BigInt(id)));
            console.log(`[Encrypt] Token array size: ${bigArr.length} tokens`);
            
            console.log('[Encrypt] Performing encryption...');
            const result = encrypt_serialize_u64_radix_flat_wasm(bigArr, clientKey);
            console.log(`[Encrypt] Encryption completed: ${result.length} bytes`);
            
            self.postMessage({ type: 'success', result });
        } catch (error) {
            console.error('[Encrypt] Encryption failed:', error.message);
            self.postMessage({ type: 'error', error: error.message });
        }
    }
};