image: Production!

Signed-off-by: Hadad <[email protected]>

assets/plugins/imageGenerator.js

@@ -6,17 +6,383 @@
 (function () {
   'use strict';
 
-  var isGenerating = false;
-  var requestId = '';
-
+  var ws = null;
+  var sessionId = '';
+  var reconnectTimer = null;
+  var isConnecting = false;
+  var currentImages = [];
+  
   if (document && document.body) {
-    isGenerating = document.body.dataset.isGenerating === 'true';
-    requestId = document.body.dataset.requestId || '';
-  }
-
-  if (isGenerating) {
-    setTimeout(function () {
-      window.location.href = '/?rid=' + encodeURIComponent(requestId);
-    }, 10000);
+    sessionId = document.body.dataset.sessionId || '';
   }
+  
+  var connectWebSocket = function() {
+    if (isConnecting || (ws && ws.readyState === 1)) return;
+    
+    isConnecting = true;
+    var protocol = window.location.protocol === 'https:' 
+      ? 'wss:' : 'ws:';
+    var wsUrl = protocol + '//' + 
+      window.location.host;
+    
+    ws = new WebSocket(wsUrl);
+    
+    ws.onopen = function() {
+      isConnecting = false;
+      if (reconnectTimer) {
+        clearTimeout(reconnectTimer);
+        reconnectTimer = null;
+      }
+      
+      ws.send(JSON.stringify({
+        type: 'register',
+        sessionId: sessionId
+      }));
+    };
+    
+    ws.onmessage = function(event) {
+      try {
+        var data = JSON.parse(event.data);
+        
+        if (!data.sessionId || 
+            data.sessionId !== sessionId) {
+          return;
+        }
+        
+        handleWebSocketMessage(data);
+      } catch (e) {}
+    };
+    
+    ws.onclose = function() {
+      isConnecting = false;
+      reconnectWebSocket();
+    };
+    
+    ws.onerror = function() {
+      isConnecting = false;
+      if (ws) ws.close();
+    };
+  };
+  
+  var reconnectWebSocket = function() {
+    if (reconnectTimer) return;
+    
+    reconnectTimer = setTimeout(function() {
+      reconnectTimer = null;
+      connectWebSocket();
+    }, 1000);
+  };
+  
+  var handleWebSocketMessage = function(data) {
+    if (!data || !data.type) return;
+    
+    switch(data.type) {
+      case 'progressUpdate':
+        updateProgressUI(data.progress);
+        break;
+        
+      case 'generationStarted':
+        showGeneratingUI();
+        break;
+        
+      case 'generationComplete':
+        handleGenerationComplete(data.images);
+        break;
+        
+      case 'generationError':
+        handleGenerationError(data.error);
+        break;
+        
+      case 'generationCancelled':
+        handleGenerationCancelled();
+        break;
+        
+      case 'imageDeleted':
+        handleImageDeleted(data.images);
+        break;
+    }
+  };
+  
+  var updateProgressUI = function(progress) {
+    var progressFill = document.querySelector('.progress-fill');
+    var progressText = document.querySelector('.progress-text');
+    
+    if (progressFill) {
+      progressFill.style.width = progress + '%';
+    }
+    
+    if (progressText) {
+      progressText.textContent = 
+        Math.floor(progress) + '% Complete';
+    }
+  };
+  
+  var showGeneratingUI = function() {
+    var outputSection = document.querySelector(
+      '.image-output-section'
+    );
+    var form = document.getElementById('generateForm');
+    var inputs = form ? 
+      form.querySelectorAll('input, select, textarea') : [];
+    
+    Array.prototype.forEach.call(inputs, function(input) {
+      input.disabled = true;
+    });
+    
+    if (outputSection) {
+      outputSection.classList.remove('has-images');
+      outputSection.innerHTML = 
+        '<div class="loading-container">' +
+        '<div class="loading-spinner" ' +
+        'style="margin: 0 auto 20px;"></div>' +
+        '<p class="loading-text">Generating your image...</p>' +
+        '<div class="progress-bar">' +
+        '<div class="progress-fill" style="width: 0%;"></div>' +
+        '</div>' +
+        '<p class="progress-text">0% Complete</p>' +
+        '</div>';
+    }
+    
+    updateButtonsForGeneration(true);
+  };
+  
+  var hideGeneratingUI = function() {
+    var form = document.getElementById('generateForm');
+    var inputs = form ? 
+      form.querySelectorAll('input, select, textarea') : [];
+    
+    Array.prototype.forEach.call(inputs, function(input) {
+      input.disabled = false;
+    });
+    
+    updateButtonsForGeneration(false);
+    window.validateInputs && window.validateInputs();
+  };
+  
+  var resetToInitialState = function() {
+    hideGeneratingUI();
+    
+    if (currentImages && currentImages.length > 0) {
+      displayImages(currentImages);
+    } else {
+      showPlaceholder();
+    }
+  };
+  
+  var showPlaceholder = function() {
+    var outputSection = document.querySelector(
+      '.image-output-section'
+    );
+    
+    if (!outputSection) return;
+    
+    outputSection.classList.remove('has-images');
+    outputSection.innerHTML = 
+      '<svg class="placeholder-icon" width="80" height="80" ' +
+      'viewBox="0 0 24 24" fill="none">' +
+      '<path d="M21 3H3C2 3 1 4 1 5V19C1 20 2 21 3 21H21C22 ' +
+      '21 23 20 23 19V5C23 4 22 3 21 3ZM21 19H3V5H21V19Z" ' +
+      'fill="currentColor"/>' +
+      '<path d="M4.5 16.5L9 12L11.5 14.5L16 10L19.5 13.5" ' +
+      'stroke="currentColor" stroke-width="1.5" ' +
+      'stroke-linecap="round"/>' +
+      '<circle cx="8" cy="8.5" r="1.5" fill="currentColor"/>' +
+      '</svg>' +
+      '<p class="placeholder-text">' +
+      'No images generated yet. Start creating amazing visuals!' +
+      '</p>';
+  };
+  
+  var updateButtonsForGeneration = function(isGenerating) {
+    var buttonsContainer = document.querySelector(
+      '.flex.justify-center.gap-4'
+    );
+    
+    if (!buttonsContainer) return;
+    
+    if (isGenerating) {
+      buttonsContainer.innerHTML = 
+        '<button type="button" onclick="cancelGeneration()" ' +
+        'class="btn btn-danger">' +
+        '<svg class="button-icon" viewBox="0 0 24 24" ' +
+        'fill="none">' +
+        '<rect x="4" y="4" width="16" height="16" rx="3" ' +
+        'fill="currentColor"/>' +
+        '</svg>' +
+        'Stop Generation' +
+        '</button>';
+    } else {
+      buttonsContainer.innerHTML = 
+        '<button type="submit" id="submitBtn" disabled ' +
+        'class="btn btn-primary">' +
+        '<svg class="button-icon" viewBox="0 0 24 24" ' +
+        'fill="none">' +
+        '<path d="M3 20V4L22 12L3 20ZM5 17L16.85 12L5 7V10.5' +
+        'L11 12L5 13.5V17Z" fill="currentColor"/>' +
+        '</svg>' +
+        'Generate Image' +
+        '</button>';
+    }
+  };
+  
+  var handleGenerationComplete = function(images) {
+    currentImages = images || [];
+    hideGeneratingUI();
+    displayImages(currentImages);
+  };
+  
+  var handleGenerationError = function(error) {
+    resetToInitialState();
+    showErrorModal(error);
+  };
+  
+  var handleGenerationCancelled = function() {
+    resetToInitialState();
+  };
+  
+  var handleImageDeleted = function(images) {
+    currentImages = images || [];
+    displayImages(currentImages);
+  };
+  
+  var displayImages = function(images) {
+    var outputSection = document.querySelector(
+      '.image-output-section'
+    );
+    
+    if (!outputSection) return;
+    
+    if (!images || images.length === 0) {
+      showPlaceholder();
+    } else {
+      outputSection.classList.add('has-images');
+      var html = '<div class="image-grid">';
+      
+      images.forEach(function(image, index) {
+        html += 
+          '<div class="image-card">' +
+          '<img src="data:image/png;base64,' + image.base64 + 
+          '" alt="' + image.prompt + '">' +
+          '<div class="image-actions">' +
+          '<a href="data:image/png;base64,' + image.base64 + 
+          '" download="generated-' + image.id + '.png" ' +
+          'class="action-btn">' +
+          '<svg class="action-icon" viewBox="0 0 24 24" ' +
+          'fill="none">' +
+          '<path d="M12 16L7 11L8.4 9.55L11 12.15V4H13V12.15' +
+          'L15.6 9.55L17 11L12 16Z" fill="currentColor"/>' +
+          '<path d="M4 20C3.45 20 2.98 19.8 2.59 19.41C2.2 ' +
+          '19.02 2 18.55 2 18V15H4V18H20V15H22V18C22 18.55 ' +
+          '21.8 19.02 21.41 19.41C21.02 19.8 20.55 20 20 20H4Z" ' +
+          'fill="currentColor"/>' +
+          '</svg>' +
+          '</a>' +
+          '<button type="button" onclick="deleteImage(' + 
+          index + ')" class="action-btn">' +
+          '<svg class="action-icon" viewBox="0 0 24 24" ' +
+          'fill="none">' +
+          '<path d="M18.3 5.71C17.91 5.32 17.28 5.32 16.89 ' +
+          '5.71L12 10.59L7.11 5.7C6.72 5.31 6.09 5.31 5.7 ' +
+          '5.7C5.31 6.09 5.31 6.72 5.7 7.11L10.59 12L5.7 ' +
+          '16.89C5.31 17.28 5.31 17.91 5.7 18.3C6.09 18.69 ' +
+          '6.72 18.69 7.11 18.3L12 13.41L16.89 18.3C17.28 ' +
+          '18.69 17.91 18.69 18.3 18.3C18.69 17.91 18.69 ' +
+          '17.28 18.3 16.89L13.41 12L18.3 7.11C18.68 6.73 ' +
+          '18.68 6.09 18.3 5.71Z" fill="currentColor"/>' +
+          '</svg>' +
+          '</button>' +
+          '</div>' +
+          '<div class="image-info">' +
+          '<p class="image-prompt">' + image.prompt + '</p>' +
+          '<p class="image-meta">' +
+          '<span class="image-model">' + 
+          image.model.toUpperCase() + '</span> | ' + 
+          image.size + '</p>' +
+          '</div>' +
+          '</div>';
+      });
+      
+      html += '</div>';
+      outputSection.innerHTML = html;
+    }
+  };
+  
+  var showErrorModal = function(error) {
+    var existingModal = document.getElementById('errorModal');
+    if (existingModal) existingModal.remove();
+    
+    var modal = document.createElement('div');
+    modal.id = 'errorModal';
+    modal.className = 'modal-overlay';
+    modal.innerHTML = 
+      '<div class="modal-content modal-error-content">' +
+      '<div class="modal-inner">' +
+      '<h3 class="modal-error-title">Error</h3>' +
+      '<p class="modal-error-text">' + error + '</p>' +
+      '<button onclick="closeErrorModal()" ' +
+      'class="btn btn-primary w-full">OK</button>' +
+      '</div>' +
+      '</div>';
+    
+    document.body.appendChild(modal);
+  };
+  
+  window.deleteImage = function(index) {
+    var xhr = new XMLHttpRequest();
+    xhr.open('POST', '/', true);
+    xhr.setRequestHeader(
+      'Content-Type', 
+      'application/json'
+    );
+    
+    xhr.onload = function() {
+      try {
+        var response = JSON.parse(xhr.responseText);
+        if (response.success) {
+          currentImages = response.images || [];
+          displayImages(currentImages);
+        }
+      } catch (e) {}
+    };
+    
+    xhr.send(JSON.stringify({
+      action: 'delete',
+      sessionId: sessionId,
+      imageIndex: index
+    }));
+  };
+  
+  var initializeImages = function() {
+    var outputSection = document.querySelector(
+      '.image-output-section'
+    );
+    
+    if (outputSection && outputSection.classList.contains('has-images')) {
+      var imageCards = outputSection.querySelectorAll('.image-card');
+      if (imageCards && imageCards.length > 0) {
+        currentImages = [];
+        imageCards.forEach(function(card) {
+          var img = card.querySelector('img');
+          var prompt = card.querySelector('.image-prompt');
+          var model = card.querySelector('.image-model');
+          var meta = card.querySelector('.image-meta');
+          
+          if (img && img.src && img.src.includes('base64,')) {
+            var base64 = img.src.split('base64,')[1];
+            var size = meta ? meta.textContent.split('|')[1] : '';
+            currentImages.push({
+              id: 'existing-' + Math.random().toString(36).substring(2, 15),
+              base64: base64,
+              prompt: prompt ? prompt.textContent : '',
+              model: model ? model.textContent.toLowerCase() : '',
+              size: size ? size.trim() : ''
+            });
+          }
+        });
+      }
+    }
+  };
+  
+  connectWebSocket();
+  initializeImages();
 })();

assets/plugins/loadParameter.js

@@ -0,0 +1,76 @@
+//
+// SPDX-FileCopyrightText: Hadad <[email protected]>
+// SPDX-License-Identifier: Apache-2.0
+//
+
+document.addEventListener('DOMContentLoaded', function() {
+    const modelSelect = document.getElementById('modelSelect');
+    const sizeSelect = document.getElementById('sizeSelect');
+    const examplesGrid = document.getElementById('examplesGrid');
+
+    if (typeof models !== 'undefined' && modelSelect) {
+        models.forEach(model => {
+            const option = document.createElement('option');
+            option.value = model.value;
+            option.textContent = model.label;
+            modelSelect.appendChild(option);
+        });
+    }
+
+    if (typeof sizes !== 'undefined' && sizeSelect) {
+        Object.keys(sizes).forEach(category => {
+            const optgroup = document.createElement('optgroup');
+            optgroup.label = category.charAt(0).toUpperCase() +
+                category.slice(1);
+
+            sizes[category].forEach(size => {
+                const option = document.createElement('option');
+                option.value = size.value;
+                option.textContent = size.label;
+                optgroup.appendChild(option);
+            });
+
+            sizeSelect.appendChild(optgroup);
+        });
+    }
+
+    if (typeof examples !== 'undefined' && examplesGrid) {
+        examples.forEach(example => {
+            const div = document.createElement('div');
+            div.className = 'example-card';
+            div.onclick = function() {
+                triggerExample(
+                    example.prompt,
+                    example.model,
+                    example.size
+                );
+            };
+
+            const promptP = document.createElement('p');
+            promptP.className = 'example-text';
+            promptP.textContent = example.prompt;
+
+            const metaP = document.createElement('p');
+            metaP.className = 'example-meta';
+
+            const modelSpan = document.createElement('span');
+            modelSpan.className = 'example-model';
+            modelSpan.textContent = example.modelLabel;
+
+            metaP.appendChild(modelSpan);
+            metaP.appendChild(
+                document.createTextNode(' | ' + example.sizeLabel)
+            );
+
+            if (example.note) {
+                metaP.appendChild(
+                    document.createTextNode(' | ' + example.note)
+                );
+            }
+
+            div.appendChild(promptP);
+            div.appendChild(metaP);
+            examplesGrid.appendChild(div);
+        });
+    }
+});

assets/plugins/webLoader.js

@@ -23,14 +23,9 @@
       });
     }
 
-    if (encode === false) {
-      return clean;
-    }
-
-    if (typeof he !== 'undefined') {
+    if (encode === true && typeof he !== 'undefined') {
       clean = he.encode(clean, {
-        useNamedReferences: true,
-        allowUnsafeSymbols: false
+        encodeEverything: false
       });
     }
 
@@ -118,9 +113,14 @@
       !!promptValue &&
       promptValue.length > 0;
 
-    if (typeof validator !== 'undefined') {
+    if (isValid) {
+      var hasNonWhitespace = /[^\s\t\n\r\u00A0\u1680\u2000-\u200A\u2028\u2029\u202F\u205F\u3000\uFEFF]/.test(promptValue);
+      isValid = isValid && hasNonWhitespace;
+    }
+
+    if (typeof validator !== 'undefined' && isValid) {
       isValid = isValid && 
-        validator.isLength(promptValue, { min: 1 });
+        validator.isLength(promptValue.trim(), { min: 1 });
     }
 
     submitBtn.disabled = !isValid;
@@ -144,137 +144,89 @@
       setTimeout(function () {
         var fa = document.getElementById('formAction');
         if (fa) fa.value = 'generate';
-        form.submit();
+        submitForm();
       }, 100);
     }
   };
 
-  var cancelGeneration = function () {
+  var submitForm = function() {
     var form = document.getElementById('generateForm');
-    var fa = document.getElementById('formAction');
-    var cancelBtn = event && event.target
-      ? event.target.closest('button')
-      : null;
-
-    if (cancelBtn) {
-      cancelBtn.disabled = true;
-      cancelBtn.style.opacity = '0.6';
-      cancelBtn.style.cursor = 'not-allowed';
-
-      var btnText = cancelBtn.querySelector('span') ||
-        cancelBtn.childNodes[cancelBtn.childNodes.length - 1];
-      
-      if (btnText && btnText.nodeType === 3) {
-        btnText.textContent = 'Cancelling...';
-      } else if (!btnText) {
-        var textNode = Array.from(cancelBtn.childNodes)
-          .find(function(node) {
-            return node.nodeType === 3 && 
-              node.textContent.trim();
-          });
-        if (textNode) {
-          textNode.textContent = 'Cancelling...';
-        }
-      }
-    }
-
-    if (fa) {
-      fa.value = 'cancel';
-
-      if (form) {
-        var hiddenCancel = document.createElement('input');
-        hiddenCancel.type = 'hidden';
-        hiddenCancel.name = 'forceCancel';
-        hiddenCancel.value = 'true';
-        form.appendChild(hiddenCancel);
-
-        var timestamp = document.createElement('input');
-        timestamp.type = 'hidden';
-        timestamp.name = 'cancelTime';
-        timestamp.value = Date.now();
-        form.appendChild(timestamp);
-
-        try {
-          form.submit();
-        } catch (e) {
-          form.requestSubmit && form.requestSubmit();
-        }
-
-        setTimeout(function() {
-          if (form && fa) {
-            fa.value = 'cancel';
-            try {
-              form.submit();
-            } catch (e) {
-              window.location.reload();
-            }
-          }
-        }, 500);
-
-        setTimeout(function() {
-          window.location.reload();
-        }, 3000);
-      }
-    } else if (form) {
-      var newFA = document.createElement('input');
-      newFA.type = 'hidden';
-      newFA.id = 'formAction';
-      newFA.name = 'action';
-      newFA.value = 'cancel';
-      form.appendChild(newFA);
-
-      try {
-        form.submit();
-      } catch (e) {
-        window.location.reload();
-      }
-    } else {
-      var xhr = new XMLHttpRequest();
-      xhr.open('POST', window.location.href, true);
-      xhr.setRequestHeader(
-        'Content-Type',
-        'application/x-www-form-urlencoded'
-      );
-      xhr.onload = function() {
-        window.location.reload();
-      };
-      xhr.onerror = function() {
-        window.location.reload();
-      };
-      xhr.send('action=cancel&forceCancel=true');
-    }
-
-    window.cancelRequested = true;
-
-    if (window.EventSource) {
-      try {
-        var sources = window.eventSources || [];
-        sources.forEach(function(source) {
-          if (source && source.close) {
-            source.close();
-          }
-        });
-      } catch (e) {}
+    if (!form) return;
+    
+    var formData = new FormData(form);
+    var promptValue = sanitizeInput(formData.get('prompt'));
+    
+    var hasNonWhitespace = /[^\s\t\n\r\u00A0\u1680\u2000-\u200A\u2028\u2029\u202F\u205F\u3000\uFEFF]/.test(promptValue);
+    
+    if (!promptValue || !hasNonWhitespace) {
+      return;
     }
-
-    if (window.AbortController && window.abortController) {
+    
+    var data = {
+      action: formData.get('action') || 'generate',
+      prompt: promptValue,
+      model: formData.get('model'),
+      size: formData.get('size'),
+      sessionId: document.body.dataset.sessionId || ''
+    };
+    
+    var xhr = new XMLHttpRequest();
+    xhr.open('POST', '/', true);
+    xhr.setRequestHeader(
+      'Content-Type', 
+      'application/json'
+    );
+    
+    xhr.onload = function() {
       try {
-        window.abortController.abort();
+        var response = JSON.parse(xhr.responseText);
+        if (!response.success && response.error) {
+          showErrorModal(response.error);
+        }
       } catch (e) {}
-    }
+    };
+    
+    xhr.send(JSON.stringify(data));
+  };
 
-    setTimeout(function() {
-      if (window.cancelRequested) {
-        window.stop && window.stop();
-        document.execCommand && 
-          document.execCommand('Stop');
-      }
-    }, 100);
+  var cancelGeneration = function () {
+    var xhr = new XMLHttpRequest();
+    xhr.open('POST', '/', true);
+    xhr.setRequestHeader(
+      'Content-Type', 
+      'application/json'
+    );
+    
+    xhr.send(JSON.stringify({
+      action: 'cancel',
+      sessionId: document.body.dataset.sessionId || ''
+    }));
+  };
+  
+  var showErrorModal = function(error) {
+    var existingModal = document.getElementById('errorModal');
+    if (existingModal) existingModal.remove();
+    
+    var modal = document.createElement('div');
+    modal.id = 'errorModal';
+    modal.className = 'modal-overlay';
+    modal.innerHTML = 
+      '<div class="modal-content modal-error-content">' +
+      '<div class="modal-inner">' +
+      '<h3 class="modal-error-title">Error</h3>' +
+      '<p class="modal-error-text">' + error + '</p>' +
+      '<button onclick="closeErrorModal()" ' +
+      'class="btn btn-primary w-full">OK</button>' +
+      '</div>' +
+      '</div>';
+    
+    document.body.appendChild(modal);
   };
 
   var modelSelect = document.getElementById('modelSelect');
   var sizeSelect = document.getElementById('sizeSelect');
   var promptInput = document.getElementById('promptInput');
+  var form = document.getElementById('generateForm');
 
   if (modelSelect) {
     modelSelect.addEventListener('change', validateInputs);
@@ -302,17 +254,40 @@
         e.preventDefault();
         var submitBtn = document.getElementById('submitBtn');
         if (submitBtn && !submitBtn.disabled) {
-          var form = document.getElementById('generateForm');
-          if (form) form.submit();
+          submitForm();
         }
       }
     });
+    
+    promptInput.addEventListener('keyup', function () {
+      validateInputs();
+    });
+    
+    promptInput.addEventListener('blur', function () {
+      validateInputs();
+    });
+  }
+  
+  if (form) {
+    form.addEventListener('submit', function(e) {
+      e.preventDefault();
+      
+      var promptValue = sanitizeInput(promptInput ? promptInput.value : '');
+      var hasNonWhitespace = /[^\s\t\n\r\u00A0\u1680\u2000-\u200A\u2028\u2029\u202F\u205F\u3000\uFEFF]/.test(promptValue);
+      
+      if (!promptValue || !hasNonWhitespace) {
+        return;
+      }
+      
+      submitForm();
+    });
   }
 
   window.closeWelcomeModal = closeWelcomeModal;
   window.closeErrorModal = closeErrorModal;
   window.triggerExample = triggerExample;
   window.cancelGeneration = cancelGeneration;
+  window.validateInputs = validateInputs;
 
   validateInputs();
 
@@ -329,11 +304,13 @@
 
   document.addEventListener('DOMContentLoaded', function() {
     InitWelcomeModal();
+    validateInputs();
   });
 
   if (document.readyState === 'complete' || 
       document.readyState === 'interactive') {
     InitWelcomeModal();
+    validateInputs();
   }
 
   setInterval(function () {

config.js

@@ -23,11 +23,15 @@ export default {
   },
   generation: {
     progressInterval: 800,
-    startDelay: 100,
+    startDelay: 50,
     maxProgress: 99
   },
   paths: {
     views: 'public',
     mainView: 'webViewer'
+  },
+  websocket: {
+    heartbeat: 30000,
+    reconnectDelay: 100
   }
 };

example.js

@@ -0,0 +1,43 @@
+//
+// SPDX-FileCopyrightText: Hadad <[email protected]>
+// SPDX-License-Identifier: Apache-2.0
+//
+
+const examples = [
+    {
+        prompt: 'A majestic mountain landscape at golden hour with dramatic clouds and vibrant colors reflecting on a crystal clear alpine lake',
+        model: 'flux',
+        size: '1536x1024',
+        modelLabel: 'FLUX',
+        sizeLabel: '1536×1024',
+        note: 'Recommended model'
+    },
+    {
+        prompt: 'An enchanted forest with bioluminescent plants and magical floating particles, moonlight streaming through ancient mystical trees',
+        model: 'kontext',
+        size: '1024x1024',
+        modelLabel: 'KONTEXT',
+        sizeLabel: '1024×1024',
+        note: 'Sometimes the server is busy'
+    },
+    {
+        prompt: 'Underwater coral reef teeming with colorful tropical fish and sea life, sunbeams penetrating crystal clear turquoise water',
+        model: 'turbo',
+        size: '1024x768',
+        modelLabel: 'TURBO',
+        sizeLabel: '1024×768',
+        note: 'Sometimes the server is busy'
+    },
+    {
+        prompt: 'A woman riding a horse across an open field, cinematic and realistic style',
+        model: 'nanobanana',
+        size: '1024x1024',
+        modelLabel: 'NANO BANANA',
+        sizeLabel: '1024×1024',
+        note: 'Sometimes the resolution is forced to be square'
+    }
+];
+
+if (typeof module !== 'undefined' && module.exports) {
+    module.exports = { examples };
+}

model.js

@@ -0,0 +1,15 @@
+//
+// SPDX-FileCopyrightText: Hadad <[email protected]>
+// SPDX-License-Identifier: Apache-2.0
+//
+
+const models = [
+    { value: 'flux', label: 'Flux' },
+    { value: 'kontext', label: 'Kontext' },
+    { value: 'turbo', label: 'Turbo' },
+    { value: 'nanobanana', label: 'Nano Banana' }
+];
+
+if (typeof module !== 'undefined' && module.exports) {
+    module.exports = { models };
+}

package.json

@@ -1,6 +1,6 @@
 {
   "name": "Image Generation Playground",
-  "version": "0.0.2",
+  "version": "1.0.0",
   "description": "Part of the UltimaX Intelligence ecosystem",
   "type": "module",
   "main": "server.js",
@@ -19,7 +19,8 @@
   "dependencies": {
     "axios": "latest",
     "ejs": "latest",
-    "express": "latest"
+    "express": "latest",
+    "ws": "latest"
   },
   "homepage": "https://umint-image.hf.space"
 }

public/webViewer.ejs

@@ -68,12 +68,17 @@
           type="image/x-icon" />
 
     <title>Image Generation Playground</title>
+
     <script src="https://cdn.tailwindcss.com"></script>
 
-    <link href="https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700;800;900&family=JetBrains+Mono:wght@400;500;600&display=swap" rel="stylesheet">
+    <link href="https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700;800;900&family=JetBrains+Mono:wght@400;500;600&display=swap" 
+          rel="stylesheet">
 
     <link rel="stylesheet" 
-          href="https://cdnjs.cloudflare.com/ajax/libs/animate.css/4.1.1/animate.min.css">
+          href="https://cdnjs.cloudflare.com/ajax/libs/animate.css/4.1.1/animate.css" 
+          integrity="sha512-phGxLIsvHFArdI7IyLjv14dchvbVkEDaH95efvAae/y2exeWBQCQDpNFbOTdV1p4/pIa/XtbuDCnfhDEIXhvGQ==" 
+          crossorigin="anonymous" 
+          referrerpolicy="no-referrer" />
 
     <link rel="stylesheet" href="/__public__/assets/css/styles.css" />
     <link rel="stylesheet" href="/__public__/assets/css/webLoader.css" />
@@ -82,13 +87,17 @@
     <script src="https://cdnjs.cloudflare.com/ajax/libs/validator/13.15.15/validator.js"></script>
     <script src="https://cdnjs.cloudflare.com/ajax/libs/he/1.2.0/he.js"></script>
     <script src="https://cdn.jsdelivr.net/npm/[email protected]/dist/xss.js"></script>
+
+    <script src="/__public__/data/models.js"></script>
+    <script src="/__public__/data/resolution.js"></script>
+    <script src="/__public__/data/examples.js"></script>
 </head>
 <body data-is-generating="<%= !!isGenerating %>"
-      data-request-id="<%= requestId || '' %>">
+      data-session-id="<%= sessionId || '' %>">
     <div class="gradient-bg"></div>
     <div class="gradient-mesh"></div>
     
-        <div id="welcomeModal" class="modal-overlay" style="display: none;">
+    <div id="welcomeModal" class="modal-overlay" style="display: none;">
         <div class="modal-content">
             <div class="modal-inner">
                 <h2 class="header-gradient animate__animated 
@@ -188,18 +197,15 @@
                                           fill="currentColor"/>
                                 </svg>
                             </a>
-                            <form method="POST" action="/">
-                                <input type="hidden" name="action" value="delete">
-                                <input type="hidden" name="imageIndex" 
-                                       value="<%= index %>">
-                                <button type="submit" class="action-btn">
-                                    <svg class="action-icon" viewBox="0 0 24 24" 
-                                         fill="none">
-                                        <path d="M18.3 5.71C17.91 5.32 17.28 5.32 16.89 5.71L12 10.59L7.11 5.7C6.72 5.31 6.09 5.31 5.7 5.7C5.31 6.09 5.31 6.72 5.7 7.11L10.59 12L5.7 16.89C5.31 17.28 5.31 17.91 5.7 18.3C6.09 18.69 6.72 18.69 7.11 18.3L12 13.41L16.89 18.3C17.28 18.69 17.91 18.69 18.3 18.3C18.69 17.91 18.69 17.28 18.3 16.89L13.41 12L18.3 7.11C18.68 6.73 18.68 6.09 18.3 5.71Z" 
-                                              fill="currentColor"/>
-                                    </svg>
-                                </button>
-                            </form>
+                            <button type="button" 
+                                    onclick="deleteImage(<%= index %>)" 
+                                    class="action-btn">
+                                <svg class="action-icon" viewBox="0 0 24 24" 
+                                     fill="none">
+                                    <path d="M18.3 5.71C17.91 5.32 17.28 5.32 16.89 5.71L12 10.59L7.11 5.7C6.72 5.31 6.09 5.31 5.7 5.7C5.31 6.09 5.31 6.72 5.7 7.11L10.59 12L5.7 16.89C5.31 17.28 5.31 17.91 5.7 18.3C6.09 18.69 6.72 18.69 7.11 18.3L12 13.41L16.89 18.3C17.28 18.69 17.91 18.69 18.3 18.3C18.69 17.91 18.69 17.28 18.3 16.89L13.41 12L18.3 7.11C18.68 6.73 18.68 6.09 18.3 5.71Z" 
+                                          fill="currentColor"/>
+                                </svg>
+                            </button>
                         </div>
                         <div class="image-info">
                             <p class="image-prompt"><%= image.prompt %></p>
@@ -234,6 +240,8 @@
             <form method="POST" action="/" id="generateForm">
                 <input type="hidden" name="action" value="generate" 
                        id="formAction">
+                <input type="hidden" name="sessionId" 
+                       value="<%= sessionId %>" />
                 
                 <div class="grid md:grid-cols-2 gap-6 form-grid">
                     <div>
@@ -242,10 +250,6 @@
                                 class="input-field select-field" 
                                 <%= isGenerating ? 'disabled' : '' %>>
                             <option value="">Select a model</option>
-                            <option value="flux">Flux</option>
-                            <option value="kontext">Kontext</option>
-                            <option value="turbo">Turbo</option>
-                            <option value="nanobanana">Nano Banana</option>
                         </select>
                     </div>
                     
@@ -255,30 +259,6 @@
                                 class="input-field select-field" 
                                 <%= isGenerating ? 'disabled' : '' %>>
                             <option value="">Select size</option>
-                            <optgroup label="Square">
-                                <option value="256x256">256×256</option>
-                                <option value="512x512">512×512</option>
-                                <option value="768x768">768×768</option>
-                                <option value="1024x1024">1024×1024</option>
-                                <option value="1536x1536">1536×1536</option>
-                                <option value="2048x2048">2048×2048</option>
-                            </optgroup>
-                            <optgroup label="Portrait">
-                                <option value="256x384">256×384</option>
-                                <option value="384x512">384×512</option>
-                                <option value="512x768">512×768</option>
-                                <option value="768x1024">768×1024</option>
-                                <option value="1024x1536">1024×1536</option>
-                                <option value="1536x2048">1536×2048</option>
-                            </optgroup>
-                            <optgroup label="Landscape">
-                                <option value="384x256">384×256</option>
-                                <option value="512x384">512×384</option>
-                                <option value="768x512">768×512</option>
-                                <option value="1024x768">1024×768</option>
-                                <option value="1536x1024">1536×1024</option>
-                                <option value="2048x1536">2048×1536</option>
-                            </optgroup>
                         </select>
                     </div>
                 </div>
@@ -321,54 +301,8 @@
         <section class="card animate__animated animate__fadeIn 
                         animate__delay-1s examples-section">
             <h3 class="examples-title">Example Prompts</h3>
-            <div class="grid lg:grid-cols-3 md:grid-cols-2 gap-4">
-                <div class="example-card" 
-                     onclick="triggerExample('A majestic mountain landscape at golden hour with dramatic clouds and vibrant colors reflecting on a crystal clear alpine lake', 'flux', '1536x1024')">
-                    <p class="example-text">
-                        A majestic mountain landscape at golden hour with 
-                        dramatic clouds and vibrant colors reflecting on a 
-                        crystal clear alpine lake
-                    </p>
-                    <p class="example-meta">
-                        <span class="example-model">FLUX</span> | 1536×1024
-                    </p>
-                </div>
-                
-                <div class="example-card" 
-                     onclick="triggerExample('An enchanted forest with bioluminescent plants and magical floating particles, moonlight streaming through ancient mystical trees', 'kontext', '1024x1024')">
-                    <p class="example-text">
-                        An enchanted forest with bioluminescent plants and 
-                        magical floating particles, moonlight streaming 
-                        through ancient mystical trees
-                    </p>
-                    <p class="example-meta">
-                        <span class="example-model">KONTEXT</span> | 1024×1024 | Sometimes the server is busy
-                    </p>
-                </div>
-                
-                <div class="example-card" 
-                     onclick="triggerExample('Underwater coral reef teeming with colorful tropical fish and sea life, sunbeams penetrating crystal clear turquoise water', 'turbo', '1024x768')">
-                    <p class="example-text">
-                        Underwater coral reef teeming with colorful tropical 
-                        fish and sea life, sunbeams penetrating crystal 
-                        clear turquoise water
-                    </p>
-                    <p class="example-meta">
-                        <span class="example-model">TURBO</span> | 1024×768 | Sometimes the server is busy
-                    </p>
-                </div>
-                
-                <div class="example-card" 
-                     onclick="triggerExample('A woman riding a horse across an open field, cinematic and realistic style', 'nanobanana', '1024x768')">
-                    <p class="example-text">
-                        A woman riding a horse across an open field, 
-                        cinematic and realistic style
-                    </p>
-                    <p class="example-meta">
-                        <span class="example-model">NANO BANANA</span> | 1024×768
-                    </p>
-                </div>
-            </div>
+            <div class="grid lg:grid-cols-3 md:grid-cols-2 gap-4" 
+                 id="examplesGrid"></div>
         </section>
     </main>
     
@@ -386,6 +320,7 @@
             </a>
         </div>
     </footer>
+    <script src="/__public__/assets/plugins/loadParameter.js"></script>
     <script defer src="/__public__/assets/plugins/imageGenerator.js"></script>
     <script defer src="/__public__/assets/plugins/webLoader.js"></script>
 </body>

resolution.js

@@ -0,0 +1,35 @@
+//
+// SPDX-FileCopyrightText: Hadad <[email protected]>
+// SPDX-License-Identifier: Apache-2.0
+//
+
+const sizes = {
+    square: [
+        { value: '256x256', label: '256×256' },
+        { value: '512x512', label: '512×512' },
+        { value: '768x768', label: '768×768' },
+        { value: '1024x1024', label: '1024×1024' },
+        { value: '1536x1536', label: '1536×1536' },
+        { value: '2048x2048', label: '2048×2048' }
+    ],
+    portrait: [
+        { value: '256x384', label: '256×384' },
+        { value: '384x512', label: '384×512' },
+        { value: '512x768', label: '512×768' },
+        { value: '768x1024', label: '768×1024' },
+        { value: '1024x1536', label: '1024×1536' },
+        { value: '1536x2048', label: '1536×2048' }
+    ],
+    landscape: [
+        { value: '384x256', label: '384×256' },
+        { value: '512x384', label: '512×384' },
+        { value: '768x512', label: '768×512' },
+        { value: '1024x768', label: '1024×768' },
+        { value: '1536x1024', label: '1536×1024' },
+        { value: '2048x1536', label: '2048×1536' }
+    ]
+};
+
+if (typeof module !== 'undefined' && module.exports) {
+    module.exports = { sizes };
+}

server.js

@@ -4,6 +4,8 @@
 //
 
 import express from 'express';
+import { createServer } from 'http';
+import { WebSocketServer } from 'ws';
 import { fileURLToPath } from 'url';
 import path from 'path';
 import config from './config.js';
@@ -13,23 +15,43 @@ import { initCleanup } from
   './src/services/storageManager.js';
 import { setupViewEngine } from 
   './src/middleware/viewEngine.js';
+import { setupWebSocket } from 
+  './src/services/websocketManager.js';
 
 const __filename = fileURLToPath(import.meta.url);
 const __dirname = path.dirname(__filename);
-
 const app = express();
+const mapping = {
+  "/__public__/assets": "assets",
+  "/__public__/data/models.js": "model.js",
+  "/__public__/data/resolution.js": "resolution.js",
+  "/__public__/data/examples.js": "example.js"
+};
+const server = createServer(app);
+const wss = new WebSocketServer({ server });
 
 setupViewEngine(app, __dirname);
+setupWebSocket(wss);
 
 app.use(express.urlencoded({ 
   extended: true, 
   limit: config.limits.bodySize 
 }));
 
-app.use(
-  "/__public__/assets",
-  express.static(path.resolve("assets"))
-);
+app.use(express.json({ 
+  limit: config.limits.bodySize 
+}));
+
+for (const [route, target] 
+  of Object.entries(mapping)
+) {
+  app.use(
+    route, 
+    express.static(
+      path.resolve(target)
+    )
+  );
+}
 
 app.use(express.static(
   path.join(__dirname, 'public')
@@ -39,7 +61,7 @@ app.use('/', imageRoutes);
 
 initCleanup();
 
-app.listen(
+server.listen(
   config.server.port, 
   config.server.host
 );

src/controllers/imageController.js

@@ -20,19 +20,22 @@ import {
   cancelGeneration 
 } from '../services/imageGenerator.js';
 import { renderView } from '../utils/viewRenderer.js';
+import { 
+  sendToSession 
+} from '../services/websocketManager.js';
 
 export const renderHome = (req, res) => {
-  const requestId = req.query.rid || generateId();
-  const data = getStorage(requestId) || getDefaultData();
+  const sessionId = generateId();
+  const data = getDefaultData();
   
   renderView(res, {
     ...data,
-    requestId
+    sessionId
   });
 };
 
 export const handleAction = async (req, res) => {
-  const requestId = req.body.requestId || generateId();
+  const sessionId = req.body.sessionId || generateId();
   const { 
     action, 
     prompt, 
@@ -41,23 +44,23 @@ export const handleAction = async (req, res) => {
     imageIndex 
   } = req.body;
   
-  let data = getStorage(requestId) || getDefaultData();
+  let data = getStorage(sessionId) || getDefaultData();
   
   switch(action) {
     case 'delete':
       return handleDelete(
-        requestId, 
+        sessionId, 
         imageIndex, 
         data, 
         res
       );
       
     case 'cancel':
-      return handleCancel(requestId, res);
+      return handleCancel(sessionId, res);
       
     case 'generate':
       return handleGenerate(
-        requestId, 
+        sessionId, 
         prompt, 
         model, 
         size, 
@@ -66,12 +69,15 @@ export const handleAction = async (req, res) => {
       );
       
     default:
-      return res.redirect(`/?rid=${requestId}`);
+      return res.json({ 
+        success: false, 
+        error: 'Invalid action' 
+      });
   }
 };
 
 const handleDelete = (
-  requestId, 
+  sessionId, 
   imageIndex, 
   data, 
   res
@@ -81,18 +87,34 @@ const handleDelete = (
     data.images.length
   )) {
     data.images.splice(parseInt(imageIndex), 1);
-    setStorage(requestId, data);
+    setStorage(sessionId, data);
+    
+    sendToSession(sessionId, {
+      type: 'imageDeleted',
+      images: data.images
+    });
   }
-  return res.redirect(`/?rid=${requestId}`);
+  
+  return res.json({ 
+    success: true, 
+    images: data.images 
+  });
 };
 
-const handleCancel = (requestId, res) => {
-  cancelGeneration(requestId);
-  return res.redirect(`/?rid=${requestId}`);
+const handleCancel = (sessionId, res) => {
+  cancelGeneration(sessionId);
+  
+  sendToSession(sessionId, {
+    type: 'generationCancelled'
+  });
+  
+  return res.json({ 
+    success: true 
+  });
 };
 
 const handleGenerate = async (
-  requestId, 
+  sessionId, 
   prompt, 
   model, 
   size, 
@@ -106,33 +128,37 @@ const handleGenerate = async (
     model, 
     size
   )) {
-    data.error = 'Please fill in all required fields';
-    setStorage(requestId, data);
-    return renderView(res, {
-      ...data,
-      requestId
+    return res.json({ 
+      success: false, 
+      error: 'Please fill in all required fields' 
     });
   }
   
   if (!validateApiConfig()) {
-    data.error = 'The server is currently busy. ' +
-                'Please try again later.';
-    setStorage(requestId, data);
-    return renderView(res, {
-      ...data,
-      requestId
+    return res.json({ 
+      success: false, 
+      error: 'The server is currently busy. ' +
+             'Please try again later.'
     });
   }
   
   data.isGenerating = true;
-  data.progress = 50;
+  data.progress = 0;
   data.error = null;
-  setStorage(requestId, data);
+  setStorage(sessionId, data);
   
-  res.redirect(`/?rid=${requestId}`);
+  sendToSession(sessionId, {
+    type: 'generationStarted',
+    progress: 0
+  });
+  
+  res.json({ 
+    success: true, 
+    sessionId 
+  });
   
   await generateImage(
-    requestId, 
+    sessionId, 
     trimmedPrompt, 
     model, 
     size

src/services/imageGenerator.js

@@ -13,18 +13,26 @@ import {
   setActiveGeneration,
   deleteActiveGeneration 
 } from './storageManager.js';
+import { 
+  sendToSession 
+} from './websocketManager.js';
 
-const updateProgress = (requestId, progress) => {
-  const data = getStorage(requestId);
+const updateProgress = (sessionId, progress) => {
+  const data = getStorage(sessionId);
   if (data) {
     data.progress = progress;
-    setStorage(requestId, data);
+    setStorage(sessionId, data);
+    
+    sendToSession(sessionId, {
+      type: 'progressUpdate',
+      progress
+    });
   }
 };
 
-const createProgressUpdater = (requestId) => {
+const createProgressUpdater = (sessionId) => {
   return setInterval(() => {
-    const data = getStorage(requestId);
+    const data = getStorage(sessionId);
     if (data && 
         data.isGenerating && 
         data.progress < config.generation.maxProgress) {
@@ -33,7 +41,7 @@ const createProgressUpdater = (requestId) => {
         config.generation.maxProgress,
         data.progress + increment
       );
-      updateProgress(requestId, newProgress);
+      updateProgress(sessionId, newProgress);
     }
   }, config.generation.progressInterval);
 };
@@ -80,15 +88,15 @@ const callImageApi = async (
 };
 
 export const generateImage = async (
-  requestId, 
+  sessionId, 
   prompt, 
   model, 
   size
 ) => {
   const controller = new AbortController();
-  setActiveGeneration(requestId, controller);
+  setActiveGeneration(sessionId, controller);
   
-  const progressInterval = createProgressUpdater(requestId);
+  const progressInterval = createProgressUpdater(sessionId);
   
   setTimeout(async () => {
     try {
@@ -99,10 +107,10 @@ export const generateImage = async (
         controller.signal
       );
       
-      const data = getStorage(requestId);
+      const data = getStorage(sessionId);
       if (!data) return;
       
-      updateProgress(requestId, 100);
+      updateProgress(sessionId, 100);
       
       if (response.data?.data?.length > 0) {
         const base64 = response.data.data[0].b64_json;
@@ -118,40 +126,50 @@ export const generateImage = async (
       
       data.isGenerating = false;
       data.progress = 0;
-      setStorage(requestId, data);
+      setStorage(sessionId, data);
+      
+      sendToSession(sessionId, {
+        type: 'generationComplete',
+        images: data.images
+      });
       
     } catch (error) {
-      const data = getStorage(requestId);
+      const data = getStorage(sessionId);
       if (!data) return;
       
       if (error.name !== 'CanceledError' && 
           error.code !== 'ERR_CANCELED') {
         data.error = 'The server is currently busy. ' +
                      'Please try again later.';
+        
+        sendToSession(sessionId, {
+          type: 'generationError',
+          error: data.error
+        });
       }
       
       data.isGenerating = false;
       data.progress = 0;
-      setStorage(requestId, data);
+      setStorage(sessionId, data);
       
     } finally {
       clearInterval(progressInterval);
-      deleteActiveGeneration(requestId);
+      deleteActiveGeneration(sessionId);
     }
   }, config.generation.startDelay);
 };
 
-export const cancelGeneration = (requestId) => {
-  const controller = getActiveGeneration(requestId);
+export const cancelGeneration = (sessionId) => {
+  const controller = getActiveGeneration(sessionId);
   if (controller) {
     controller.abort();
-    deleteActiveGeneration(requestId);
+    deleteActiveGeneration(sessionId);
     
-    const data = getStorage(requestId);
+    const data = getStorage(sessionId);
     if (data) {
       data.isGenerating = false;
       data.progress = 0;
-      setStorage(requestId, data);
+      setStorage(sessionId, data);
     }
   }
 };

src/services/storageManager.js

@@ -8,37 +8,37 @@ import config from '../../config.js';
 const tempStorage = new Map();
 const activeGenerations = new Map();
 
-export const getStorage = (requestId) => {
-  const data = tempStorage.get(requestId);
+export const getStorage = (sessionId) => {
+  const data = tempStorage.get(sessionId);
   if (data) {
     data.lastAccess = Date.now();
   }
   return data;
 };
 
-export const setStorage = (requestId, data) => {
+export const setStorage = (sessionId, data) => {
   data.lastAccess = Date.now();
-  tempStorage.set(requestId, data);
+  tempStorage.set(sessionId, data);
 };
 
-export const deleteStorage = (requestId) => {
-  tempStorage.delete(requestId);
-  activeGenerations.delete(requestId);
+export const deleteStorage = (sessionId) => {
+  tempStorage.delete(sessionId);
+  activeGenerations.delete(sessionId);
 };
 
-export const getActiveGeneration = (requestId) => {
-  return activeGenerations.get(requestId);
+export const getActiveGeneration = (sessionId) => {
+  return activeGenerations.get(sessionId);
 };
 
 export const setActiveGeneration = (
-  requestId, 
+  sessionId, 
   controller
 ) => {
-  activeGenerations.set(requestId, controller);
+  activeGenerations.set(sessionId, controller);
 };
 
-export const deleteActiveGeneration = (requestId) => {
-  activeGenerations.delete(requestId);
+export const deleteActiveGeneration = (sessionId) => {
+  activeGenerations.delete(sessionId);
 };
 
 export const getDefaultData = () => ({

src/services/websocketManager.js

@@ -0,0 +1,108 @@
+//
+// SPDX-FileCopyrightText: Hadad <[email protected]>
+// SPDX-License-Identifier: Apache-2.0
+//
+
+import config from '../../config.js';
+import { generateId } from '../utils/idGenerator.js';
+
+const clients = new Map();
+
+export const setupWebSocket = (wss) => {
+  wss.on('connection', (ws) => {
+    const clientId = generateId();
+    let clientSessionId = null;
+    
+    ws.clientId = clientId;
+    ws.isAlive = true;
+    
+    ws.on('pong', () => {
+      ws.isAlive = true;
+    });
+    
+    ws.on('message', (data) => {
+      try {
+        const message = JSON.parse(data.toString());
+        
+        if (message.type === 'register' && message.sessionId) {
+          clientSessionId = message.sessionId;
+          
+          if (clients.has(clientSessionId)) {
+            const oldClient = clients.get(clientSessionId);
+            if (oldClient && oldClient.readyState === 1) {
+              oldClient.close();
+            }
+          }
+          
+          clients.set(clientSessionId, ws);
+          ws.sessionId = clientSessionId;
+          
+          ws.send(JSON.stringify({
+            type: 'registered',
+            sessionId: clientSessionId
+          }));
+        }
+        
+        handleWebSocketMessage(ws, message);
+      } catch (error) {}
+    });
+    
+    ws.on('close', () => {
+      if (clientSessionId) {
+        clients.delete(clientSessionId);
+      }
+    });
+    
+    ws.on('error', () => {
+      if (clientSessionId) {
+        clients.delete(clientSessionId);
+      }
+    });
+    
+    ws.send(JSON.stringify({
+      type: 'connected',
+      clientId
+    }));
+  });
+  
+  const interval = setInterval(() => {
+    wss.clients.forEach((ws) => {
+      if (ws.isAlive === false) {
+        return ws.terminate();
+      }
+      ws.isAlive = false;
+      ws.ping();
+    });
+  }, config.websocket.heartbeat);
+  
+  wss.on('close', () => {
+    clearInterval(interval);
+  });
+};
+
+const handleWebSocketMessage = (ws, message) => {
+  if (message.type === 'ping') {
+    ws.send(JSON.stringify({ type: 'pong' }));
+  }
+};
+
+export const sendToSession = (sessionId, data) => {
+  const client = clients.get(sessionId);
+  if (client && client.readyState === 1) {
+    client.send(JSON.stringify({
+      ...data,
+      sessionId
+    }));
+  }
+};
+
+export const broadcastToAll = (data) => {
+  clients.forEach((client, sessionId) => {
+    if (client.readyState === 1) {
+      client.send(JSON.stringify({
+        ...data,
+        sessionId
+      }));
+    }
+  });
+};