From 2e8fa6e30b6526cdac71a1a75a8752fe07d3ec41 Mon Sep 17 00:00:00 2001
From: Russell Ballestrini <russell.ballestrini@gmail.com>
Date: Tue, 1 Jul 2025 19:44:12 -0400
Subject: [PATCH] fix(translation): resolve preview issues and endpoint errors

- Remove modal/dialog DOM from translation preview content
- Fix 404 errors by properly checking model registry for endpoints
- Enable floating button in preview windows with partial initialization
---
 src/models.js      | 22 +++++++++++++++++-----
 src/translation.js |  9 ++++++++-
 src/ui.js          | 11 ++++++++---
 3 files changed, 33 insertions(+), 9 deletions(-)

diff --git a/src/models.js b/src/models.js
index f5b4ef8..059d812 100644
--- a/src/models.js
+++ b/src/models.js
@@ -221,13 +221,25 @@ export function getSelectedModel() {
 
 // Helper function to get the API endpoint for the selected model.
 export function getSelectedModelEndpoint() {
-	const selectedModel = getSelectedModel();
-
-	if (modelRegistry[selectedModel]) {
-		return modelRegistry[selectedModel].url;
+	// Check modal dropdown first
+	const modalDropdown = document.getElementById("modal-model-selection");
+	if (modalDropdown && modalDropdown.value && modelRegistry[modalDropdown.value]) {
+		return modelRegistry[modalDropdown.value].url;
 	}
 
-	// Fallback to the first endpoint if model not found
+	// Fallback to main dropdown
+	const dropdown = document.getElementById("model-selection");
+	if (dropdown && dropdown.value && modelRegistry[dropdown.value]) {
+		return modelRegistry[dropdown.value].url;
+	}
+
+	// Find the first available model in the registry
+	const availableModels = Object.keys(modelRegistry);
+	if (availableModels.length > 0) {
+		return modelRegistry[availableModels[0]].url;
+	}
+
+	// Fallback to the first endpoint if model registry is empty
 	if (VLLM_ENDPOINTS.length > 0) {
 		return VLLM_ENDPOINTS[0].url;
 	}
diff --git a/src/translation.js b/src/translation.js
index 2321260..f8b0c8a 100644
--- a/src/translation.js
+++ b/src/translation.js
@@ -264,8 +264,15 @@ ${preservedText}`;
 
 // Extract page content for translation
 export function extractPageContent() {
+	// Clone the document to avoid modifying the original
+	const documentClone = document.cloneNode(true);
+	
+	// Remove any open modals/dialogs that shouldn't be in the translation
+	const modalsToRemove = documentClone.querySelectorAll('dialog[open], #hermes-modal, [id*="modal"], [class*="modal"]');
+	modalsToRemove.forEach(modal => modal.remove());
+	
 	// Return the entire HTML of the document to preserve head, styles, and scripts
-	return document.documentElement.outerHTML;
+	return documentClone.documentElement.outerHTML;
 }
 
 // Translate current page content
diff --git a/src/ui.js b/src/ui.js
index 76dd98f..403549a 100644
--- a/src/ui.js
+++ b/src/ui.js
@@ -1883,13 +1883,18 @@ export function openTranslateModal() {
 
 // Initialize the system
 export function initializeSystem() {
-	// Prevent duplicate initialization in popup windows
+	// Check for skip init flag - allow partial initialization for preview windows
 	if (window.UNCLOSEAI_SKIP_INIT === true) {
-		console.log("uncloseai.js: Skipping initialization as requested by flag.");
+		console.log("uncloseai.js: Partial initialization for preview window");
+		// Only create floating button in preview windows if not explicitly disabled
+		if (window.UNCLOSEAI_FLOATING_BUTTON !== false) {
+			console.log("uncloseai.js: Creating floating button in preview window");
+			createFloatingAIButton();
+		}
 		return;
 	}
 
-	console.log("uncloseai.js: Initializing system");
+	console.log("uncloseai.js: Full initialization");
 	initializeChatInterface();
 
 	// Only create floating button if not disabled