From 593cdbe6ccccb606928430b2ae027b6ca5568842 Mon Sep 17 00:00:00 2001
From: Kieran Klukas <me@dunkirk.sh>
Date: Wed, 28 Jan 2026 22:49:25 -0500
Subject: [PATCH] feat: add token validation, file size tracking, and polling
 support
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Validate Hack Club AI token format (sk-hc-v1-<64 hex chars>)
- Show file size before/after with percentage change
- Display dimensions and current size in info bar
- Add polling infrastructure for async background removal
- Falls back to synchronous mode if proxy doesn't support async
- Show elapsed time during processing
- Move status messages below image preview

💘 Generated with Crush

Assisted-by: Claude Sonnet 4.5 via Crush <crush@charm.land>
---
 public/index.html | 208 ++++++++++++++++++++++++++++++++++++++++------
 src/index.js      | 104 ++++++++++++++++++++---
 2 files changed, 275 insertions(+), 37 deletions(-)

diff --git a/public/index.html b/public/index.html
index 9c3a1f0..a6c0580 100644
--- a/public/index.html
+++ b/public/index.html
@@ -112,6 +112,16 @@
             padding: 1rem;
             color: #666;
         }
+        .file-info {
+            text-align: center;
+            padding: 0.5rem;
+            color: #666;
+            font-size: 0.9rem;
+            background: #f9f9f9;
+            border-radius: 4px;
+            margin-bottom: 1rem;
+        }
+        .file-info strong { color: #333; }
         .spinner {
             display: inline-block;
             width: 20px;
@@ -240,9 +250,7 @@
             </div>
         </div>
 
-        <div id="statusArea" class="status hidden">
-            <span class="spinner"></span><span id="statusText">Processing...</span>
-        </div>
+        <div id="fileInfo" class="file-info hidden"></div>
 
         <!-- Normal preview -->
         <div id="previewArea" class="preview-area">
@@ -272,6 +280,10 @@
                 </div>
             </div>
         </div>
+
+        <div id="statusArea" class="status hidden">
+            <span class="spinner"></span><span id="statusText">Processing...</span>
+        </div>
     </div>
 
     <script>
@@ -286,6 +298,7 @@
         let dragStart = { x: 0, y: 0 };
         let cropStart = { x: 0, y: 0, w: 0, h: 0 };
         let inCropMode = false;
+        let currentFileSize = 0;
 
         // Elements
         const tokenRow = document.getElementById('tokenRow');
@@ -304,6 +317,7 @@
         const bgColor = document.getElementById('bgColor');
         const showTransparent = document.getElementById('showTransparent');
         const aspectRatio = document.getElementById('aspectRatio');
+        const fileInfo = document.getElementById('fileInfo');
 
         // Token handling
         if (localStorage.getItem('hc_token')) {
@@ -311,8 +325,18 @@
         }
 
         function saveToken() {
-            localStorage.setItem('hc_token', tokenInput.value);
+            const token = tokenInput.value.trim();
+            
+            // Validate token format: sk-hc-v1-<64 hex chars>
+            const tokenRegex = /^sk-hc-v1-[a-f0-9]{64}$/i;
+            if (!tokenRegex.test(token)) {
+                showError('Invalid token format. Expected: sk-hc-v1-<64 hex characters>');
+                return;
+            }
+            
+            localStorage.setItem('hc_token', token);
             tokenRow.classList.add('hidden');
+            hideError();
         }
 
         function showError(msg) {
@@ -354,6 +378,7 @@
         function loadFile(f) {
             if (!f) return;
             hideError();
+            currentFileSize = f.size;
             const reader = new FileReader();
             reader.onload = e => {
                 const img = new Image();
@@ -362,6 +387,7 @@
                     dropZone.classList.add('hidden');
                     editor.classList.remove('hidden');
                     updatePreview();
+                    updateFileInfo();
                 };
                 img.src = e.target.result;
             };
@@ -409,6 +435,42 @@
             }
 
             ctx.drawImage(currentImage, 0, 0);
+            updateFileInfo();
+        }
+
+        function formatFileSize(bytes) {
+            if (bytes < 1024) return bytes + ' B';
+            if (bytes < 1024 * 1024) return (bytes / 1024).toFixed(1) + ' KB';
+            return (bytes / (1024 * 1024)).toFixed(2) + ' MB';
+        }
+
+        function getCanvasSize() {
+            return new Promise(resolve => {
+                previewCanvas.toBlob(blob => resolve(blob.size), 'image/png');
+            });
+        }
+
+        async function updateFileInfo() {
+            if (!currentImage) {
+                fileInfo.classList.add('hidden');
+                return;
+            }
+
+            const currentSize = await getCanvasSize();
+            const dims = `${currentImage.width}×${currentImage.height}px`;
+            
+            let info = `<strong>Dimensions:</strong> ${dims} | <strong>Current size:</strong> ${formatFileSize(currentSize)}`;
+            
+            if (currentFileSize && currentFileSize !== currentSize) {
+                const diff = currentSize - currentFileSize;
+                const percent = ((diff / currentFileSize) * 100).toFixed(1);
+                const color = diff > 0 ? '#c00' : '#0a0';
+                info += ` | <strong>Original:</strong> ${formatFileSize(currentFileSize)} `;
+                info += `<span style="color: ${color}">(${diff > 0 ? '+' : ''}${percent}%)</span>`;
+            }
+            
+            fileInfo.innerHTML = info;
+            fileInfo.classList.remove('hidden');
         }
 
         function getCanvasDataUrl() {
@@ -460,8 +522,12 @@
             tempCanvas.getContext('2d').drawImage(currentImage, sx, sy, sw, sh, 0, 0, sw, sh);
 
             const img = new Image();
-            img.onload = () => {
+            img.onload = async () => {
                 currentImage = img;
+                // Update file size after crop
+                currentFileSize = await new Promise(resolve => {
+                    tempCanvas.toBlob(blob => resolve(blob.size), 'image/png');
+                });
                 inCropMode = false;
                 cropArea.classList.add('hidden');
                 previewArea.classList.remove('hidden');
@@ -618,14 +684,22 @@
             }
         }
 
-        // Remove background
+        // Remove background with polling
         async function removeBg() {
-            const tk = localStorage.getItem('hc_token') || tokenInput.value;
+            const tk = localStorage.getItem('hc_token') || tokenInput.value.trim();
             if (!tk) { showError('Enter token first'); return; }
+            
+            // Validate token format
+            const tokenRegex = /^sk-hc-v1-[a-f0-9]{64}$/i;
+            if (!tokenRegex.test(tk)) {
+                showError('Invalid token format. Expected: sk-hc-v1-<64 hex characters>');
+                return;
+            }
+            
             if (!currentImage) return;
 
             hideError();
-            showStatus('Removing background...');
+            showStatus('Starting background removal...');
             document.getElementById('removeBgBtn').disabled = true;
 
             // Get current image as data URL
@@ -639,6 +713,7 @@
             const selectedModel = document.getElementById('bgRemoverModel').value;
 
             try {
+                // Start the prediction
                 const res = await fetch('/api/remove-bg', {
                     method: 'POST',
                     headers: { 'Content-Type': 'application/json' },
@@ -647,20 +722,34 @@
                 const data = await res.json();
                 if (!res.ok) throw new Error(typeof data.error === 'string' ? data.error : JSON.stringify(data.error));
 
-                const img = new Image();
-                img.crossOrigin = 'anonymous';
-                img.onload = () => {
-                    currentImage = img;
-                    hideStatus();
-                    document.getElementById('removeBgBtn').disabled = false;
-                    updatePreview();
-                };
-                img.onerror = () => {
-                    hideStatus();
-                    document.getElementById('removeBgBtn').disabled = false;
-                    showError('Failed to load result image');
-                };
-                img.src = data.output;
+                // Check if we got direct output (no polling needed) or prediction ID
+                if (data.predictionId) {
+                    // Poll for status
+                    await pollStatus(data.predictionId, tk);
+                } else if (data.status === 'succeeded' && data.output) {
+                    // Direct output, load immediately
+                    const img = new Image();
+                    img.crossOrigin = 'anonymous';
+                    img.onload = async () => {
+                        currentImage = img;
+                        currentFileSize = await new Promise(resolve => {
+                            const tempCanvas = document.createElement('canvas');
+                            tempCanvas.width = img.width;
+                            tempCanvas.height = img.height;
+                            tempCanvas.getContext('2d').drawImage(img, 0, 0);
+                            tempCanvas.toBlob(blob => resolve(blob.size), 'image/png');
+                        });
+                        hideStatus();
+                        document.getElementById('removeBgBtn').disabled = false;
+                        updatePreview();
+                    };
+                    img.onerror = () => {
+                        hideStatus();
+                        document.getElementById('removeBgBtn').disabled = false;
+                        showError('Failed to load result image');
+                    };
+                    img.src = data.output;
+                }
             } catch (e) {
                 hideStatus();
                 document.getElementById('removeBgBtn').disabled = false;
@@ -668,6 +757,67 @@
             }
         }
 
+        async function pollStatus(predictionId, token) {
+            const startTime = Date.now();
+            
+            const poll = async () => {
+                try {
+                    const res = await fetch(`/api/check-status?id=${predictionId}&token=${encodeURIComponent(token)}`);
+                    const data = await res.json();
+                    
+                    if (!res.ok) throw new Error(typeof data.error === 'string' ? data.error : JSON.stringify(data.error));
+
+                    const elapsed = Math.floor((Date.now() - startTime) / 1000);
+                    
+                    if (data.status === 'succeeded') {
+                        showStatus(`Processing complete (${elapsed}s)`);
+                        
+                        const img = new Image();
+                        img.crossOrigin = 'anonymous';
+                        img.onload = async () => {
+                            currentImage = img;
+                            // Reset currentFileSize to the new processed image size
+                            currentFileSize = await new Promise(resolve => {
+                                const tempCanvas = document.createElement('canvas');
+                                tempCanvas.width = img.width;
+                                tempCanvas.height = img.height;
+                                tempCanvas.getContext('2d').drawImage(img, 0, 0);
+                                tempCanvas.toBlob(blob => resolve(blob.size), 'image/png');
+                            });
+                            hideStatus();
+                            document.getElementById('removeBgBtn').disabled = false;
+                            updatePreview();
+                        };
+                        img.onerror = () => {
+                            hideStatus();
+                            document.getElementById('removeBgBtn').disabled = false;
+                            showError('Failed to load result image');
+                        };
+                        img.src = data.output;
+                    } else if (data.status === 'failed') {
+                        hideStatus();
+                        document.getElementById('removeBgBtn').disabled = false;
+                        showError(data.error || 'Background removal failed');
+                    } else if (data.status === 'processing' || data.status === 'starting') {
+                        showStatus(`Processing... ${elapsed}s elapsed`);
+                        // Poll again in 2 seconds
+                        setTimeout(poll, 2000);
+                    } else {
+                        // Unknown status, keep polling
+                        showStatus(`Status: ${data.status} (${elapsed}s)`);
+                        setTimeout(poll, 2000);
+                    }
+                } catch (e) {
+                    hideStatus();
+                    document.getElementById('removeBgBtn').disabled = false;
+                    showError(e.message);
+                }
+            };
+
+            // Start polling
+            poll();
+        }
+
         // Apply background color (bakes it into the image)
         function applyBgColor() {
             if (!currentImage) return;
@@ -682,8 +832,12 @@
             ctx.drawImage(currentImage, 0, 0);
 
             const img = new Image();
-            img.onload = () => {
+            img.onload = async () => {
                 currentImage = img;
+                // Update file size after applying bg color
+                currentFileSize = await new Promise(resolve => {
+                    tempCanvas.toBlob(blob => resolve(blob.size), 'image/png');
+                });
                 updatePreview();
             };
             img.src = tempCanvas.toDataURL('image/png');
@@ -710,6 +864,7 @@
         // Reset
         function resetAll() {
             currentImage = null;
+            currentFileSize = 0;
             inCropMode = false;
             editor.classList.add('hidden');
             dropZone.classList.remove('hidden');
@@ -717,6 +872,7 @@
             previewArea.classList.remove('hidden');
             hideError();
             hideStatus();
+            fileInfo.classList.add('hidden');
             fileInput.value = '';
         }
 
@@ -759,8 +915,12 @@
 
             // Load back as current image
             const img = new Image();
-            img.onload = () => {
+            img.onload = async () => {
                 currentImage = img;
+                // Update file size after optimization
+                currentFileSize = await new Promise(resolve => {
+                    tempCanvas.toBlob(blob => resolve(blob.size), 'image/jpeg', quality);
+                });
                 updatePreview();
             };
             img.src = optimized;
diff --git a/src/index.js b/src/index.js
index 4fb6594..6a9aae0 100644
--- a/src/index.js
+++ b/src/index.js
@@ -16,6 +16,10 @@ export default {
       return handleRemoveBg(request, env);
     }
 
+    if (url.pathname === "/api/check-status" && request.method === "GET") {
+      return handleCheckStatus(request, env);
+    }
+
     // For non-API routes, return 404 to let assets handle it
     return new Response("Not Found", { status: 404 });
   },
@@ -40,32 +44,106 @@ async function handleRemoveBg(request, env) {
 
     const input = { image };
 
-    // Select model based on parameter
+    // Select model based on parameter - use full model ID format
     const modelId = model === '851labs' 
       ? "851-labs/background-remover:a029dff38972b5fda4ec5d75d7d1cd25aeff621d2cf4946a41055d7db66b80bc"
       : "lucataco/remove-bg:95fcc2a26d3899cd6c2691c900465aaeff466285a65c14638cc5f36f34befaf1";
 
-    const output = await replicate.run(modelId, { input });
+    console.log("Running model:", modelId);
+
+    // Use replicate.run with stream option to get prediction ID
+    const prediction = await replicate.run(modelId, { 
+      input,
+      wait: false  // Don't wait for completion
+    });
+
+    console.log("Prediction response:", prediction);
 
-    // Handle different output formats
-    let outputUrl;
-    if (typeof output === "string") {
-      outputUrl = output;
-    } else if (output && typeof output.url === "function") {
-      outputUrl = output.url();
-    } else if (output && output.url) {
-      outputUrl = output.url;
+    // Check if we got a prediction object or direct output
+    if (prediction && prediction.id) {
+      // Got a prediction object, can poll it
+      return jsonResponse({ 
+        predictionId: prediction.id,
+        status: prediction.status 
+      });
     } else {
-      outputUrl = output;
+      // Got direct output (proxy may have waited), return it
+      let outputUrl;
+      if (typeof prediction === "string") {
+        outputUrl = prediction;
+      } else if (Array.isArray(prediction) && prediction.length > 0) {
+        outputUrl = prediction[0];
+      } else if (prediction && prediction.url) {
+        outputUrl = prediction.url;
+      } else {
+        outputUrl = prediction;
+      }
+      
+      return jsonResponse({ 
+        status: 'succeeded',
+        output: outputUrl 
+      });
     }
-
-    return jsonResponse({ output: outputUrl });
   } catch (e) {
     console.error("Error:", e);
     return jsonResponse({ error: e.message }, 500);
   }
 }
 
+async function handleCheckStatus(request, env) {
+  try {
+    const url = new URL(request.url);
+    const predictionId = url.searchParams.get('id');
+    const token = url.searchParams.get('token');
+
+    if (!predictionId) {
+      return jsonResponse({ error: "Missing prediction ID" }, 400);
+    }
+
+    if (!token) {
+      return jsonResponse({ error: "Missing token" }, 400);
+    }
+
+    const replicate = new Replicate({
+      auth: token,
+      baseUrl: "https://ai.hackclub.com/proxy/v1/replicate",
+    });
+
+    const prediction = await replicate.predictions.get(predictionId);
+
+    // Return status and output if available
+    const response = {
+      status: prediction.status,
+    };
+
+    if (prediction.status === 'succeeded' && prediction.output) {
+      // Handle different output formats
+      let outputUrl;
+      if (typeof prediction.output === "string") {
+        outputUrl = prediction.output;
+      } else if (prediction.output && typeof prediction.output.url === "function") {
+        outputUrl = prediction.output.url();
+      } else if (prediction.output && prediction.output.url) {
+        outputUrl = prediction.output.url;
+      } else if (Array.isArray(prediction.output) && prediction.output.length > 0) {
+        outputUrl = prediction.output[0];
+      } else {
+        outputUrl = prediction.output;
+      }
+      response.output = outputUrl;
+    }
+
+    if (prediction.status === 'failed') {
+      response.error = prediction.error || 'Prediction failed';
+    }
+
+    return jsonResponse(response);
+  } catch (e) {
+    console.error("Error checking status:", e);
+    return jsonResponse({ error: e.message }, 500);
+  }
+}
+
 function corsHeaders() {
   return {
     "Access-Control-Allow-Origin": "*",
-- 
2.51.2