compatible with server

wasm-demo.js

@@ -2,7 +2,7 @@ import initWasm, {
     decrypt_serialized_u64_radix_flat_wasm
 } from './concrete-ml-extensions-wasm/concrete_ml_extensions_wasm.js';
 
-const SERVER = 'https://backend-server.com';
+const SERVER = 'https://api.zama.ai';
 
 let clientKey, serverKey;
 let encTokens;
@@ -10,6 +10,7 @@ let encServerResult;
 let keygenWorker;
 let encryptWorker;
 let sessionUid;
+let taskId;
 
 // Memory-efficient base64 encoding for large Uint8Array
 function uint8ToBase64(uint8) {
@@ -78,37 +79,43 @@ show('tokenizerSpin', false);
                     // Initialize the worker with the client key
                     encryptWorker.postMessage({ type: 'init', clientKey });
                     
-                    console.log('[Main] Encoding server key to base64...');
-                    const serverKeyB64 = await uint8ToBase64(serverKey);
-                    console.log(`[Main] Server key base64 size: ${serverKeyB64.length} bytes`);
-                    const handshakeResponse = await fetch(`${SERVER}/handshake`, {
-                        method:'POST', 
-                        headers:{'Content-Type':'application/json'},
-                        body: JSON.stringify({
-                            server_key_b64: serverKeyB64
-                        })
+                    console.log('[Main] Sending server key to server...');
+                    $('keygenStatus').textContent = 'Keys generated, sending server key...';
+                    show('keygenSpin', true);
+
+                    const formData = new FormData();
+                    const serverKeyBlob = new Blob([serverKey], { type: 'application/octet-stream' });
+                    const serverKeyFile = new File([serverKeyBlob], "server.key");
+                    formData.append('key', serverKeyFile);
+                    formData.append('task_name', 'synthid');
+
+                    const addKeyResponse = await fetch(`${SERVER}/add_key`, {
+                        method: 'POST',
+                        body: formData
                     });
-                    
-                    if (!handshakeResponse.ok) {
-                        const errorText = await handshakeResponse.text();
-                        throw new Error(`Server handshake failed: ${handshakeResponse.status} ${errorText}`);
+
+                    if (!addKeyResponse.ok) {
+                        const errorText = await addKeyResponse.text();
+                        throw new Error(`Server /add_key failed: ${addKeyResponse.status} ${errorText}`);
                     }
-                    
-                    const { uid } = await handshakeResponse.json();
+
+                    const { uid } = await addKeyResponse.json();
                     sessionUid = uid;
-                    console.log('[Main] Server handshake successful');
-                    $('keygenStatus').textContent = 'keys generated ✓';
+                    console.log('[Main] Server key sent and UID received:', sessionUid);
+                    $('keygenStatus').textContent = 'Keys generated & UID received ✓';
                     enable('btnEncrypt');
                 } catch (error) {
-                    console.error('[Main] Server handshake error:', error);
-                    $('keygenStatus').textContent = `Server handshake failed: ${error.message}`;
+                    console.error('[Main] Server key submission error:', error);
+                    $('keygenStatus').textContent = `Server key submission failed: ${error.message}`;
                     enable('btnEncrypt', false);
+                } finally {
+                    show('keygenSpin', false);
                 }
             } else {
                 console.error('[Main] Key generation error:', e.data.error);
                 $('keygenStatus').textContent = `Error generating keys: ${e.data.error}`;
+                show('keygenSpin', false);
             }
-            show('keygenSpin', false);
         };
     } catch (e) {
         console.error('[Main] Failed to initialize WASM module:', e);
@@ -164,79 +171,125 @@ $('btnEncrypt').onclick = async () => {
     }
 };
 
+async function pollTaskStatus(currentTaskId, currentUid) {
+    try {
+        const statusResponse = await fetch(`${SERVER}/get_task_status?task_id=${currentTaskId}&uid=${currentUid}`);
+        if (!statusResponse.ok) {
+            const errorText = await statusResponse.text();
+            console.error(`[Poll] Error fetching status: ${statusResponse.status} ${errorText}`);
+            $('srvStatus').textContent = `Status check error: ${statusResponse.status}`;
+            show('spin', false);
+            return null;
+        }
+
+        const statusData = await statusResponse.json();
+        console.log('[Poll] Task status:', statusData);
+        $('srvStatus').textContent = `Status: ${statusData.status} - ${statusData.details}`;
+
+        if (statusData.status === 'success' || statusData.status === 'completed') {
+            return statusData;
+        } else if (['failure', 'revoked', 'unknown', 'error'].includes(statusData.status.toLowerCase())) {
+            console.error('[Poll] Task failed or unrecoverable:', statusData);
+            $('srvStatus').textContent = `Task failed: ${statusData.status}`;
+            show('spin', false);
+            return null;
+        } else {
+            setTimeout(() => pollTaskStatus(currentTaskId, currentUid).then(finalStatus => {
+                if (finalStatus && (finalStatus.status === 'success' || finalStatus.status === 'completed')) {
+                    getTaskResult(currentTaskId, currentUid, 'synthid');
+                }
+            }), 5000);
+            return null;
+        }
+    } catch (e) {
+        console.error('[Poll] Polling exception:', e);
+        $('srvStatus').textContent = `Polling error: ${e.message}`;
+        show('spin', false);
+        return null;
+    }
+}
+
+async function getTaskResult(currentTaskId, currentUid, taskName) {
+    $('srvStatus').textContent = 'Fetching result...';
+    try {
+        const resultResponse = await fetch(`${SERVER}/get_task_result?task_name=${taskName}&task_id=${currentTaskId}&uid=${currentUid}`);
+
+        if (!resultResponse.ok) {
+            const errorText = await resultResponse.text();
+            throw new Error(`Server /get_task_result error: ${resultResponse.status} ${errorText}`);
+        }
+
+        const resultArrayBuffer = await resultResponse.arrayBuffer();
+        encServerResult = new Uint8Array(resultArrayBuffer);
+
+        console.log(`[Main] Received encrypted result: ${encServerResult.length} bytes`);
+        $('encResult').value = `(${encServerResult.length} B)`;
+        $('srvStatus').textContent = `✓ result received (${((performance.now() - window.taskStartTime) / 1000).toFixed(2)}s total)`;
+        enable('btnDecrypt');
+
+    } catch (e) {
+        const duration = window.taskStartTime ? ((performance.now() - window.taskStartTime) / 1000).toFixed(2) : 'N/A';
+        console.error(`[Main] /get_task_result failed after ${duration}s:`, e);
+        $('srvStatus').textContent = `Result fetch error: ${e.message} (${duration}s)`;
+    } finally {
+        show('spin', false);
+        $('srvComputing').hidden = true;
+    }
+}
+
 $('btnSend').onclick = async () => {
     if ($('spin').hidden === false) {
-        console.log('[Main] Send already in progress, ignoring click');
+        console.log('[Main] Task submission/polling already in progress, ignoring click');
+        return;
+    }
+    if (!sessionUid || !encTokens) {
+        alert('Please generate keys and encrypt text first.');
         return;
     }
+
     show('encIcon', false);
     show('spin', true);
-    $('srvStatus').textContent = 'Preparing request…';
+    $('srvStatus').textContent = 'Submitting task…';
     $('srvComputing').hidden = true;
-    $('srvProgress').hidden = true;
-    const startTime = performance.now();
+    window.taskStartTime = performance.now();
+
     try {
-        const ciphertext_b64 = await uint8ToBase64(encTokens);
-
-        $('srvStatus').textContent = 'Sending…';
-        console.log('[Main] Sending request to server...');
-        const response = await fetch(`${SERVER}/detect`, {
-            method:'POST', 
-            headers:{'Content-Type':'application/json'},
-            body: JSON.stringify({
-                uid: sessionUid,
-                ciphertext_b64
-            })
-        }).catch(error => {
-            throw new Error(`Server unreachable: ${error.message}`);
+        const formData = new FormData();
+        formData.append('uid', sessionUid);
+        formData.append('task_name', 'synthid');
+
+        const encryptedInputBlob = new Blob([encTokens], { type: 'application/octet-stream' });
+        const encryptedInputFile = new File([encryptedInputBlob], "input.fheencrypted");
+        formData.append('encrypted_input', encryptedInputFile);
+
+        const startTaskResponse = await fetch(`${SERVER}/start_task`, {
+            method: 'POST',
+            body: formData
         });
-        
-        console.log(`[Main] Server response status: ${response.status}`);
-        if (!response.ok) {
-            const errorText = await response.text();
-            throw new Error(`Server error: ${response.status} ${errorText}`);
+
+        if (!startTaskResponse.ok) {
+            const errorText = await startTaskResponse.text();
+            throw new Error(`Server /start_task error: ${startTaskResponse.status} ${errorText}`);
         }
 
-        $('srvStatus').textContent = '✓ sent';
+        const { task_id: newTaskId } = await startTaskResponse.json();
+        taskId = newTaskId;
+        console.log('[Main] Task submitted to server. Task ID:', taskId);
+        $('srvStatus').textContent = `Task submitted (ID: ${taskId.substring(0,8)}...). Polling status...`;
         $('srvComputing').hidden = false;
-        $('srvProgress').hidden = false;
-        
-        const tokenCount = llama3Tokenizer.encode($('tokenInput').value.trim()).length;
-        const estimatedTime = tokenCount * 15;
-        const endTime = startTime + (estimatedTime * 1000);
-        
-        const progressInterval = setInterval(() => {
-            const now = performance.now();
-            const elapsed = (now - startTime) / 1000;
-            const remaining = Math.max(0, estimatedTime - elapsed);
-            $('srvProgress').textContent = `Estimated time remaining: ${Math.ceil(remaining)}s`;
-            
-            if (remaining <= 0) {
-                clearInterval(progressInterval);
+
+        pollTaskStatus(taskId, sessionUid).then(finalStatus => {
+            if (finalStatus && (finalStatus.status === 'success' || finalStatus.status === 'completed')) {
+                getTaskResult(taskId, sessionUid, 'synthid');
             }
-        }, 1000);
-        
-        const {result_b64} = await response.json();
-        clearInterval(progressInterval);
-        const actualEndTime = performance.now();
-        const duration = ((actualEndTime - startTime) / 1000).toFixed(2);
-        console.log(`[Main] Server request completed in ${duration}s`);
-        console.log('[Main] Processing server response...');
-        encServerResult = Uint8Array.from(atob(result_b64), c=>c.charCodeAt(0));
-        console.log(`[Main] Received encrypted result: ${encServerResult.length} bytes`);
-        $('encResult').value = `(${encServerResult.length} B)`;
-        $('srvComputing').hidden = true;
-        $('srvProgress').hidden = true;
-        $('srvStatus').textContent = `✓ received (${duration}s)`;
-        enable('btnDecrypt');
+        });
+
     } catch (e) {
-        const endTime = performance.now();
-        const duration = ((endTime - startTime) / 1000).toFixed(2);
-        console.error(`[Main] Server request failed after ${duration}s:`, e);
-        $('srvComputing').hidden = true;
-        $('srvProgress').hidden = true;
-        $('srvStatus').textContent = `Server error: ${e.message} (${duration}s)`;
+        const duration = ((performance.now() - window.taskStartTime) / 1000).toFixed(2);
+        console.error(`[Main] Task submission failed after ${duration}s:`, e);
+        $('srvStatus').textContent = `Task submission error: ${e.message} (${duration}s)`;
         show('spin', false);
+        $('srvComputing').hidden = true;
     }
 };
 
@@ -248,9 +301,9 @@ $('btnDecrypt').onclick = () => {
         const score = (Number(score_scaled) / 1e6).toFixed(6);
         console.log('[Main] Decryption successful');
         console.log(`[Main] Result - flag: ${flag}, score: ${score}, total_g: ${total_g}`);
-        $('decResult').textContent = `flag=${flag}, score=${score}, total_g=${total_g}`;
+        $('decResult').textContent = `Flag: ${flag}, Score: ${score}, Total G: ${total_g}`;
     } catch (e) {
         console.error('[Main] Decryption error:', e);
-        $('decResult').textContent = `Error decrypting result: ${e.message}`;
+        $('decResult').textContent = `Decryption failed: ${e.message}`;
     }
 };