From 2fc6f929f21e83f0e7c16d7eb56dc76cd1bb76ad Mon Sep 17 00:00:00 2001
From: jordojordo <jordonleach@gmail.com>
Date: Mon, 27 May 2024 15:06:08 -0400
Subject: [PATCH] Add proxy for websocket connection

---
 .env.production                  |  5 +++-
 Containerfile                    |  2 +-
 env.d.ts                         |  4 +++
 nginx.conf.template              |  5 +++-
 scripts/start.sh                 | 18 ++++++++++---
 src/App.vue                      | 16 ++++++-----
 src/components/Config.vue        | 46 ++++++++++++++++++++++----------
 src/services/websocket.ts        | 36 ++++++++++++++++++++++---
 src/stores/useToolConfigStore.ts | 33 +++++++++++++++++++----
 src/types/chat.ts                |  8 ++++++
 src/utils/service.ts             | 15 +++++++++++
 11 files changed, 154 insertions(+), 34 deletions(-)
 create mode 100644 src/utils/service.ts

diff --git a/.env.production b/.env.production
index fd7c7dd..63aea35 100644
--- a/.env.production
+++ b/.env.production
@@ -1 +1,4 @@
-VITE_THOTHSCRIPT_API=CLIENT_THOTHSCRIPT_API
\ No newline at end of file
+VITE_WS_SCHEME=CLIENT_THOTHSCRIPT_PROXY_SCHEME
+VITE_WS_HOST=CLIENT_THOTHSCRIPT_PROXY_HOST
+VITE_WS_PORT=CLIENT_THOTHSCRIPT_PROXY_PORT
+VITE_WS_PATH=CLIENT_THOTHSCRIPT_PROXY_PATH
diff --git a/Containerfile b/Containerfile
index 907c8c9..fd42d6e 100644
--- a/Containerfile
+++ b/Containerfile
@@ -1,4 +1,4 @@
-FROM node:latest as build-stage
+FROM node:22-alpine as build-stage
 WORKDIR /app
 COPY package*.json ./
 RUN yarn
diff --git a/env.d.ts b/env.d.ts
index e26bddf..0aac3c7 100644
--- a/env.d.ts
+++ b/env.d.ts
@@ -2,6 +2,10 @@
 
 interface ImportMetaEnv {
   readonly VITE_THOTHSCRIPT_API: string;
+  readonly VITE_WS_SECURE: boolean;
+  readonly VITE_WS_HOST: string;
+  readonly VITE_WS_PORT: string;
+  readonly VITE_WS_PATH: string;
   readonly BASE_URL: string;
   readonly PROD: boolean
 }
diff --git a/nginx.conf.template b/nginx.conf.template
index e35857f..e233918 100644
--- a/nginx.conf.template
+++ b/nginx.conf.template
@@ -31,11 +31,14 @@ http {
     }
 
     location /ws/ {
-      proxy_pass __CLIENT_THOTHSCRIPT_API__;
+      proxy_pass http://__CLIENT_THOTHSCRIPT_OPERATOR_HOST__.__CLIENT_THOTHSCRIPT_OPERATOR_NAMESPACE__.svc.cluster.local:__CLIENT_THOTHSCRIPT_OPERATOR_PORT__;
       proxy_http_version 1.1;
       proxy_set_header Upgrade $http_upgrade;
       proxy_set_header Connection "upgrade";
       proxy_set_header Host $host;
+      proxy_set_header X-Real-IP $remote_addr;
+      proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
+      proxy_set_header X-Forwarded-Proto $scheme;
       access_log /var/log/nginx/websocket_access.log main;
     }
 
diff --git a/scripts/start.sh b/scripts/start.sh
index 712e2c2..e5eea6f 100644
--- a/scripts/start.sh
+++ b/scripts/start.sh
@@ -1,10 +1,22 @@
 #!/bin/sh
 
-# Provide default values for CLIENT_THOTHSCRIPT_API if it's not set
-: "${CLIENT_THOTHSCRIPT_API:=localhost:3000}"
+
+# Provide default values for CLIENT_THOTHSCRIPT_API_SCHEME, CLIENT_THOTHSCRIPT_API_HOST, CLIENT_THOTHSCRIPT_API_PORT, and CLIENT_THOTHSCRIPT_API_PATH if not set
+: "${CLIENT_THOTHSCRIPT_PROXY_SCHEME:=ws}"
+: "${CLIENT_THOTHSCRIPT_PROXY_HOST:=localhost}"
+: "${CLIENT_THOTHSCRIPT_PROXY_PORT:=80}"
+: "${CLIENT_THOTHSCRIPT_PROXY_PATH:=/ws/}"
+
+# Provide default values for CLIENT_THOTHSCRIPT_OPERATOR_HOST, CLIENT_THOTHSCRIPT_OPERATOR_NAMESPACE, and CLIENT_THOTHSCRIPT_OPERATOR_PORT if not set
+: "${CLIENT_THOTHSCRIPT_OPERATOR_HOST:=thothscript-operator}"
+: "${CLIENT_THOTHSCRIPT_OPERATOR_NAMESPACE:=thothscript}"
+: "${CLIENT_THOTHSCRIPT_OPERATOR_PORT:=3000}"
 
 # Substitute environment variables in nginx.conf.template and create nginx.conf
-sed "s|__CLIENT_THOTHSCRIPT_API__|http:\/\/$CLIENT_THOTHSCRIPT_API|g" /etc/nginx/nginx.conf.template > /etc/nginx/nginx.conf
+sed -e "s|__CLIENT_THOTHSCRIPT_OPERATOR_HOST__|${CLIENT_THOTHSCRIPT_OPERATOR_HOST}|g" \
+    -e "s|__CLIENT_THOTHSCRIPT_OPERATOR_NAMESPACE__|${CLIENT_THOTHSCRIPT_OPERATOR_NAMESPACE}|g" \
+    -e "s|__CLIENT_THOTHSCRIPT_OPERATOR_PORT__|${CLIENT_THOTHSCRIPT_OPERATOR_PORT}|g" \
+    /etc/nginx/nginx.conf.template > /etc/nginx/nginx.conf
 
 for i in $(env | grep CLIENT_); do
   key=$(echo $i | cut -d '=' -f 1)
diff --git a/src/App.vue b/src/App.vue
index da12d53..f4fcbbf 100644
--- a/src/App.vue
+++ b/src/App.vue
@@ -7,21 +7,25 @@ import websocket from '@/services/websocket';
 
 const toolConfigStore = useToolConfigStore();
 
+toolConfigStore.loadFromLocalStorage();
+
 onMounted(() => {
-  websocket.connect(toolConfigStore.websocketUrl);
+  websocket.connect();
 });
 
 onUnmounted(() => {
   websocket.disconnect();
 });
 
-// Watch for changes in the WebSocket URL and reconnect if it changes
-watch(() => toolConfigStore.websocketUrl, (newUrl, oldUrl) => {
-  if ( newUrl !== oldUrl ) {
+// Watch for changes in the WebSocket config and reconnect if it changes
+watch(() => toolConfigStore.websocket, () => {
+  if ( toolConfigStore.shouldReconnect ) {
     websocket.disconnect();
-    websocket.connect(newUrl);
+    websocket.connect();
+
+    toolConfigStore.setShouldReconnect(false);
   }
-});
+}, { deep: true });
 </script>
 
 <template>
diff --git a/src/components/Config.vue b/src/components/Config.vue
index 9389040..a35518b 100644
--- a/src/components/Config.vue
+++ b/src/components/Config.vue
@@ -1,26 +1,33 @@
 <script setup lang="ts">
 import { ref, reactive } from 'vue';
+
 import { useToolConfigStore } from '@/stores/useToolConfigStore';
 
 const toolConfigStore = useToolConfigStore();
 
 const localConfig = reactive({
-  websocketUrl: toolConfigStore.websocketUrl,
-  modelName:    toolConfigStore.modelName,
-  maxTokens:    toolConfigStore.maxTokens,
-  temperature:  toolConfigStore.temperature,
-  chat:         toolConfigStore.chat,
-  runOpts:      {
+  websocket: {
+    host: toolConfigStore.websocket.host,
+    port: toolConfigStore.websocket.port,
+    keepAliveInterval: toolConfigStore.websocket.keepAliveInterval,
+  },
+  modelName: toolConfigStore.modelName,
+  maxTokens: toolConfigStore.maxTokens,
+  temperature: toolConfigStore.temperature,
+  chat: toolConfigStore.chat,
+  runOpts: {
     disableCache: toolConfigStore.runOpts.disableCache,
-    quiet:        toolConfigStore.runOpts.quiet,
-    workspace:    toolConfigStore.runOpts.workspace,
-  }
+    quiet: toolConfigStore.runOpts.quiet,
+    workspace: toolConfigStore.runOpts.workspace,
+  },
 });
 
 const saveButton = ref('Save');
 
 const saveConfig = () => {
-  toolConfigStore.websocketUrl = localConfig.websocketUrl;
+  toolConfigStore.websocket.host = localConfig.websocket.host;
+  toolConfigStore.websocket.port = localConfig.websocket.port;
+  toolConfigStore.websocket.keepAliveInterval = localConfig.websocket.keepAliveInterval;
   toolConfigStore.modelName = localConfig.modelName;
   toolConfigStore.maxTokens = localConfig.maxTokens;
   toolConfigStore.temperature = localConfig.temperature;
@@ -29,16 +36,18 @@ const saveConfig = () => {
   toolConfigStore.runOpts.quiet = localConfig.runOpts.quiet;
   toolConfigStore.runOpts.workspace = localConfig.runOpts.workspace;
 
-  // Display feedback message
+  toolConfigStore.setShouldReconnect(true);
+  toolConfigStore.saveToLocalStorage();
+
   saveButton.value = 'Saved!';
 
-  // Hide message after 3 seconds
   setTimeout(() => {
     saveButton.value = 'Save';
   }, 3000);
 };
 </script>
 
+
 <template>
   <section class="chat-config">
     <h2>Settings</h2>
@@ -46,8 +55,16 @@ const saveConfig = () => {
       <fieldset>
         <legend>Connection</legend>
         <div>
-          <label>ThothScript Operator URL</label>
-          <input v-model="localConfig.websocketUrl" type="text" />
+          <label>Host</label>
+          <input v-model="localConfig.websocket.host" type="text" />
+        </div>
+        <div>
+          <label>Port</label>
+          <input v-model="localConfig.websocket.port" type="text" />
+        </div>
+        <div>
+          <label>Keep-Alive Interval (seconds)</label>
+          <input v-model.number="localConfig.websocket.keepAliveInterval" type="number" />
         </div>
       </fieldset>
       <fieldset>
@@ -94,6 +111,7 @@ const saveConfig = () => {
   </section>
 </template>
 
+
 <style scoped>
 .chat-config {
   display: flex;
diff --git a/src/services/websocket.ts b/src/services/websocket.ts
index a5c9610..3fa4d27 100644
--- a/src/services/websocket.ts
+++ b/src/services/websocket.ts
@@ -20,23 +20,43 @@ const state = reactive<WebSocketState>({
 
 let socket: WebSocket | null = null;
 let reconnectTimeout: NodeJS.Timeout | null = null;
+let keepAliveInterval: number | null = null;
+
 const maxReconnectAttempts = 10; // Limit number of attempts within 10 minutes
 const reconnectInterval = 10000; // Try to reconnect every 10 seconds
 
-const connect = (url?: string) => {
+const connect = () => {
   if ( reconnectTimeout !== null ) {
     clearTimeout(reconnectTimeout);
   }
 
   const toolConfigStore = useToolConfigStore();
+  
+  const KEEP_ALIVE_INTERVAL = ( toolConfigStore.websocket.keepAliveInterval || 120 ) * 1000;
+  const { secure, host, port, path } = toolConfigStore.websocket;
+
+  const scheme = secure ? 'wss' : 'ws';
+  const hostStr = host ? host : window.location.hostname;
+  const portStr = port ? `:${port}` : '';
+  const pathStr = path ? path : '/ws/';
+  const websocketUrl = `${ scheme }://${ hostStr }${ portStr }${ pathStr }`;
 
-  socket = new WebSocket(`ws://${ url || toolConfigStore.websocketUrl }/ws/`);
+  // console.log(`Connecting to WebSocket: ${websocketUrl}`);
+  socket = new WebSocket(websocketUrl);
 
   socket.onopen = () => {
     state.isConnected = true;
     state.reconnectAttempts = 0;
     emitGrowl('Connected', 'success');
     emitter.emit('connected');
+    console.log('WebSocket connected to:', websocketUrl); // eslint-disable-line no-console
+
+    // Start keep-alive interval
+    keepAliveInterval = window.setInterval(() => {
+      if ( socket && socket.readyState === WebSocket.OPEN ) {
+        socket.send(JSON.stringify({ type: 'ping' }));
+      }
+    }, KEEP_ALIVE_INTERVAL);
   };
 
   socket.onmessage = (event: MessageEvent) => {
@@ -53,10 +73,15 @@ const connect = (url?: string) => {
     if ( state.reconnectAttempts < maxReconnectAttempts ) {
       emitGrowl('Disconnected. Reconnecting...', 'warning');
       state.reconnectAttempts++;
-      reconnectTimeout = setTimeout(() => connect(url), reconnectInterval);
+      reconnectTimeout = setTimeout(() => connect(), reconnectInterval);
     } else {
       emitGrowl('Failed to reconnect after multiple attempts.', 'error');
     }
+
+    if ( keepAliveInterval ) {
+      clearInterval(keepAliveInterval);
+      keepAliveInterval = null;
+    }
   };
 
   socket.onerror = (error: Event) => {
@@ -128,6 +153,11 @@ const disconnect = () => {
   if ( socket && state.isConnected ) {
     socket.close();
   }
+
+  if ( keepAliveInterval ) {
+    clearInterval(keepAliveInterval);
+    keepAliveInterval = null;
+  }
 };
 
 const emitGrowl = (message: string, type: 'success' | 'warning' | 'error') => {
diff --git a/src/stores/useToolConfigStore.ts b/src/stores/useToolConfigStore.ts
index 8d66c49..6a93d9e 100644
--- a/src/stores/useToolConfigStore.ts
+++ b/src/stores/useToolConfigStore.ts
@@ -1,13 +1,19 @@
 import { defineStore } from 'pinia';
 import type { RunOpts } from '@gptscript-ai/gptscript';
 
+import { getWebsocketConfig } from '@/utils/service';
+import type { WebsocketConfig } from '@/types/chat';
+
+const LOCAL_STORAGE_KEY = 'toolConfig';
+
 export interface ToolConfigState {
   maxTokens: number;
   modelName: string;
   temperature: number;
   chat: boolean;
   runOpts: RunOpts;
-  websocketUrl: string;
+  websocket: WebsocketConfig;
+  shouldReconnect: boolean;
 }
 
 export const useToolConfigStore = defineStore('toolConfig', {
@@ -25,7 +31,11 @@ export const useToolConfigStore = defineStore('toolConfig', {
       subTool:      '',
       workspace:    '',
     },
-    websocketUrl: import.meta.env.VITE_THOTHSCRIPT_API || 'localhost:3000'
+    websocket: {
+      ...getWebsocketConfig(),
+      keepAliveInterval: 120,
+    },
+    shouldReconnect: false,
   }),
   actions: {
     updateMaxTokens(maxTokens: number) {
@@ -43,8 +53,21 @@ export const useToolConfigStore = defineStore('toolConfig', {
     updateRunOpts(runOpts: RunOpts) {
       this.runOpts = runOpts;
     },
-    updateWebSocketUrl(url: string) {
-      this.websocketUrl = url;
-    }
+    updateWebSocketConfig(config: WebsocketConfig) {
+      this.websocket = config;
+    },
+    setShouldReconnect(value: boolean) {
+      this.shouldReconnect = value;
+    },
+    saveToLocalStorage() {
+      localStorage.setItem(LOCAL_STORAGE_KEY, JSON.stringify(this.$state));
+    },
+    loadFromLocalStorage() {
+      const storedState = localStorage.getItem(LOCAL_STORAGE_KEY);
+
+      if ( storedState ) {
+        this.$patch(JSON.parse(storedState));
+      }
+    },
   },
 });
diff --git a/src/types/chat.ts b/src/types/chat.ts
index 91b67a5..83c14ad 100644
--- a/src/types/chat.ts
+++ b/src/types/chat.ts
@@ -18,3 +18,11 @@ export interface SystemMessage {
   output: SystemOutput
   error?: object;
 }
+
+export interface WebsocketConfig {
+  secure: boolean;
+  host: string;
+  port: string;
+  path: string;
+  keepAliveInterval?: number;
+}
diff --git a/src/utils/service.ts b/src/utils/service.ts
new file mode 100644
index 0000000..c02e6b9
--- /dev/null
+++ b/src/utils/service.ts
@@ -0,0 +1,15 @@
+import type { WebsocketConfig } from '@/types/chat';
+
+export function getWebsocketConfig(): WebsocketConfig {
+  const secure = import.meta.env.VITE_WS_SECURE === true;
+  const host = import.meta.env.VITE_WS_HOST || window.location.hostname;
+  const port = import.meta.env.VITE_WS_PORT || '';
+  const path = import.meta.env.VITE_WS_PATH || '/ws/';
+
+  return {
+    secure,
+    host,
+    port,
+    path
+  };
+}