From 9e66f42f5d0486c37ed42a5ff2d0966395e56d3d Mon Sep 17 00:00:00 2001
From: ebolam <ebolam@gmail.com>
Date: Wed, 7 Sep 2022 20:12:21 -0400
Subject: [PATCH] Cookie Saving for Colab - Needs work

---
 aiserver.py          | 21 ++++++++++++---------
 koboldai_settings.py |  1 +
 static/koboldai.js   | 45 ++++++++++++++++++++++++++++----------------
 3 files changed, 42 insertions(+), 25 deletions(-)

diff --git a/aiserver.py b/aiserver.py
index 49aabc34..50bad101 100644
--- a/aiserver.py
+++ b/aiserver.py
@@ -6776,7 +6776,7 @@ def new_ui_index():
 def ui2_connect():
     #Send all variables to client
     koboldai_vars.send_to_ui()
-    UI_2_load_tweaks()
+    UI_2_load_cookies()
     pass
     
 #==================================================================#
@@ -7692,19 +7692,22 @@ def UI_2_theme_list_refresh(data):
 #==================================================================#
 # Save Tweaks
 #==================================================================#
-@socketio.on('save_tweaks')
-def UI_2_save_tweaks(data):
-    with open("./settings/tweaks.settings", "w") as f:
-        f.write(data)
+@socketio.on('save_cookies')
+def UI_2_save_cookies(data):
+    for key in data:
+        #Note this won't sync to the client automatically as we're modifying a variable rather than setting it
+        koboldai_vars.cookies[key] = data[key]
+    with open("./settings/cookies.settings", "w") as f:
+        f.write(json.dumps(koboldai_vars.cookies))
 
 #==================================================================#
 # Load Tweaks
 #==================================================================#
-def UI_2_load_tweaks():
+def UI_2_load_cookies():
     if koboldai_vars.on_colab:
-        if os.path.exists("./settings/tweaks.settings"):
-            with open("./settings/tweaks.settings", "r") as f:
-                socketio.emit('load_tweaks', f.read(), room="UI2")
+        if os.path.exists("./settings/cookies.settings"):
+            with open("./settings/cookies.settings", "r") as f:
+                socketio.emit('load_cookies', json.load(f), room="UI2")
 
 #==================================================================#
 # Test
diff --git a/koboldai_settings.py b/koboldai_settings.py
index cbdb3c44..a985d0e7 100644
--- a/koboldai_settings.py
+++ b/koboldai_settings.py
@@ -782,6 +782,7 @@ class system_settings(settings):
         print("Colab Check: {}".format(self.on_colab))
         self.horde_share = False
         self._horde_pid = None
+        self.cookies = {} #cookies for colab since colab's URL changes, cookies are lost
         
         
     def __setattr__(self, name, value):
diff --git a/static/koboldai.js b/static/koboldai.js
index 2ce5e3d9..3d4e738e 100644
--- a/static/koboldai.js
+++ b/static/koboldai.js
@@ -25,7 +25,7 @@ socket.on("world_info_folder", function(data){world_info_folder(data);});
 socket.on("delete_new_world_info_entry", function(data){document.getElementById("world_info_-1").remove();});
 socket.on("delete_world_info_entry", function(data){document.getElementById("world_info_"+data).remove();});
 socket.on("error", function(data){show_error_message(data);});
-socket.on('load_tweaks', function(data){load_tweaks(data);});
+socket.on('load_cookies', function(data){load_cookies(data);});
 //socket.onAny(function(event_name, data) {console.log({"event": event_name, "class": data.classname, "data": data});});
 
 var presets = {};
@@ -2025,17 +2025,20 @@ function save_tweaks() {
 		let path = tweakContainer.getAttribute("tweak-path");
 		if (toggle.checked) out.push(path);
 	}
-
-	if (on_colab) {
-		socket.emit("save_tweaks", JSON.stringify(out));
-	} else {
-		setCookie("enabledTweaks", JSON.stringify(out));
-	}
+	setCookie("enabledTweaks", JSON.stringify(out));
 }
 
 
-function load_tweaks(data) {
-	let enabledTweaks = JSON.parse(data);
+function load_cookies(data) {
+	for (cookie of data) {
+		setCookie(cookie, data[cookie]);
+	}
+	process_cookies();
+}
+
+function load_tweaks() {
+	
+	let enabledTweaks = JSON.parse(getCookie("enabledTweaks", "[]"));
 
 	for (const tweakContainer of document.getElementsByClassName("tweak-container")) {
 		let toggle = tweakContainer.querySelector("input");
@@ -3093,6 +3096,9 @@ function setCookie(cname, cvalue, exdays=60) {
   const d = new Date();
   d.setTime(d.getTime() + (exdays * 24 * 60 * 60 * 1000));
   let expires = "expires="+d.toUTCString();
+  if (document.getElementById("on_colab").textContent == "true") {
+	socket.emit("save_cookies", {cname: cvavlue});
+  }
   document.cookie = cname + "=" + cvalue + ";" + expires + ";";
 }
 
@@ -3380,13 +3386,7 @@ function open_finder() {
 	finderInput.focus();
 }
 
-$(document).ready(function(){
-	on_colab = document.getElementById("on_colab").textContent == "true";
-
-	create_theming_elements();
-	document.onkeydown = detect_key_down;
-	document.onkeyup = detect_key_up;
-	document.getElementById("input_text").onkeydown = detect_enter_submit;
+function process_cookies() {
 	if (getCookie("Settings_Pin") == "false") {
 		settings_unpin();
 	} else {
@@ -3399,6 +3399,19 @@ $(document).ready(function(){
 	}
 	preserve_game_space(!(getCookie("preserve_game_space") == "false"));
 	options_on_right(!(getCookie("options_on_right") == "false"));
+	
+	load_tweaks();
+}
+
+$(document).ready(function(){
+	on_colab = document.getElementById("on_colab").textContent == "true";
+
+	create_theming_elements();
+	document.onkeydown = detect_key_down;
+	document.onkeyup = detect_key_up;
+	document.getElementById("input_text").onkeydown = detect_enter_submit;
+	
+	process_cookies();
 
 
 	// Tweak registering