From 7ad48c29adb9288e205adce4a46747eec53f71fc Mon Sep 17 00:00:00 2001
From: MaxRobinsonTheGreat <maxdrob1@gmail.com>
Date: Fri, 13 Dec 2024 10:42:41 -0600
Subject: [PATCH] added full shutdown from ui/task

---
 src/agent/agent.js           | 22 +++++++++++-----------
 src/agent/agent_proxy.js     |  4 ++++
 src/agent/tasks.js           |  2 +-
 src/process/main_proxy.js    | 10 ++++++++++
 src/server/mind_server.js    | 22 ++++++++++++++++++++++
 src/server/public/index.html | 10 ++++++++++
 6 files changed, 58 insertions(+), 12 deletions(-)

diff --git a/src/agent/agent.js b/src/agent/agent.js
index 5e9cc5f..79e1d29 100644
--- a/src/agent/agent.js
+++ b/src/agent/agent.js
@@ -438,22 +438,18 @@ export class Agent {
         }, INTERVAL);
 
         this.bot.emit('idle');
-
-        // Check for task completion
-        if (this.task.data) {
-            setInterval(() => {
-                let res = this.task.isDone();
-                if (res) {
-                    // TODO kill other bots
-                    this.cleanKill(res.message, res.code);
-                }
-            }, 1000);
-        }
     }
 
     async update(delta) {
         await this.bot.modes.update();
         this.self_prompter.update(delta);
+        if (this.task.data) {
+            let res = this.task.isDone();
+            if (res) {
+                console.log('Task finished:', res.message);
+                this.killAll();
+            }
+        }
     }
 
     isIdle() {
@@ -466,4 +462,8 @@ export class Agent {
         this.history.save();
         process.exit(code);
     }
+
+    killAll() {
+        serverProxy.shutdown();
+    }
 }
diff --git a/src/agent/agent_proxy.js b/src/agent/agent_proxy.js
index 81395ac..feeba37 100644
--- a/src/agent/agent_proxy.js
+++ b/src/agent/agent_proxy.js
@@ -48,6 +48,10 @@ class AgentServerProxy {
         this.socket.emit('login-agent', this.agent.name);
     }
 
+    shutdown() {
+        this.socket.emit('shutdown');
+    }
+
     getSocket() {
         return this.socket;
     }
diff --git a/src/agent/tasks.js b/src/agent/tasks.js
index 881b933..f7527f1 100644
--- a/src/agent/tasks.js
+++ b/src/agent/tasks.js
@@ -179,7 +179,7 @@ export class Task {
             await new Promise((resolve) => setTimeout(resolve, 10000));
             if (available_agents.length < this.data.agent_count) {
                 console.log(`Missing ${this.data.agent_count - available_agents.length} bot(s).`);
-                this.agent.cleanKill('Not all required players/bots are present in the world. Exiting.', 4);
+                this.agent.killAll();
             }
         }
 
diff --git a/src/process/main_proxy.js b/src/process/main_proxy.js
index 44c2733..8336458 100644
--- a/src/process/main_proxy.js
+++ b/src/process/main_proxy.js
@@ -35,6 +35,16 @@ class MainProxy {
         this.socket.on('register-agents-success', () => {
             console.log('Agents registered');
         });
+
+        this.socket.on('shutdown', () => {
+            console.log('Shutting down');
+            for (let agentName in this.agent_processes) {
+                this.agent_processes[agentName].stop();
+            }
+            setTimeout(() => {
+                process.exit(0);
+            }, 2000);
+        });
     }
 
     addAgent(agent) {
diff --git a/src/server/mind_server.js b/src/server/mind_server.js
index ae952d3..5d99290 100644
--- a/src/server/mind_server.js
+++ b/src/server/mind_server.js
@@ -101,6 +101,19 @@ export function createMindServer(port = 8080) {
             }
         });
 
+        socket.on('stop-all-agents', () => {
+            console.log('Killing all agents');
+            stopAllAgents();
+        });
+
+        socket.on('shutdown', () => {
+            console.log('Shutting down');
+            for (let manager of Object.values(agentManagers)) {
+                manager.emit('shutdown');
+            }
+            process.exit(0);
+        });
+
     });
 
     server.listen(port, 'localhost', () => {
@@ -121,6 +134,15 @@ function agentsUpdate(socket) {
     socket.emit('agents-update', agents);
 }
 
+function stopAllAgents() {
+    for (const agentName in inGameAgents) {
+        let manager = agentManagers[agentName];
+        if (manager) {
+            manager.emit('stop-agent', agentName);
+        }
+    }
+}
+
 // Optional: export these if you need access to them from other files
 export const getIO = () => io;
 export const getServer = () => server;
diff --git a/src/server/public/index.html b/src/server/public/index.html
index b597ed9..43ada99 100644
--- a/src/server/public/index.html
+++ b/src/server/public/index.html
@@ -85,6 +85,8 @@
                             `}
                         </div>
                     </div>
+                    <button class="stop-btn" onclick="killAllAgents()">Stop All</button>
+                    <button class="stop-btn" onclick="shutdown()">Shutdown</button>
                 `).join('') :
                 '<div class="agent">No agents connected</div>';
         });
@@ -100,6 +102,14 @@
         function stopAgent(agentName) {
             socket.emit('stop-agent', agentName);
         }
+
+        function killAllAgents() {
+            socket.emit('stop-all-agents');
+        }
+
+        function shutdown() {
+            socket.emit('shutdown');
+        }
     </script>
 </body>
 </html> 
\ No newline at end of file