From 829359584bb23614049078b097b54da830d5aa7e Mon Sep 17 00:00:00 2001
From: Michael Neale <michael.neale@gmail.com>
Date: Thu, 16 Jan 2025 14:45:39 +1100
Subject: [PATCH] chore [v1.0]: tidy up processes when session closes (#622)

---
 crates/goose-cli/src/log_usage.rs | 15 ++++++++++-----
 ui/desktop/src/goosed.ts          | 23 ++++-------------------
 ui/desktop/src/main.ts            |  7 ++++---
 3 files changed, 18 insertions(+), 27 deletions(-)

diff --git a/crates/goose-cli/src/log_usage.rs b/crates/goose-cli/src/log_usage.rs
index 8c21d2e98..862078efa 100644
--- a/crates/goose-cli/src/log_usage.rs
+++ b/crates/goose-cli/src/log_usage.rs
@@ -61,7 +61,11 @@ mod tests {
     fn test_session_logging() {
         run_with_tmp_dir(|| {
             let home_dir = dirs::home_dir().unwrap();
-            let log_dir = home_dir.join(".config").join("goose").join("logs");
+            let log_file = home_dir
+                .join(".config")
+                .join("goose")
+                .join("logs")
+                .join("goose.log");
 
             log_usage(
                 "path.txt".to_string(),
@@ -73,12 +77,10 @@ mod tests {
             );
 
             // Check if log file exists and contains the expected content
-            let log_file = log_dir.join("goose.log");
-            assert!(log_file.exists());
+            assert!(log_file.exists(), "Log file should exist");
 
             let log_content = std::fs::read_to_string(&log_file).unwrap();
-            let log: SessionLog =
-                serde_json::from_str(log_content.lines().last().unwrap()).unwrap();
+            let log: SessionLog = serde_json::from_str(&log_content).unwrap();
 
             assert!(log.session_file.contains("path.txt"));
             assert_eq!(log.usage[0].usage.input_tokens, Some(10));
@@ -86,6 +88,9 @@ mod tests {
             assert_eq!(log.usage[0].usage.total_tokens, Some(30));
             assert_eq!(log.usage[0].model, "model");
             assert_eq!(log.usage[0].cost, Some(dec!(0.5)));
+
+            // Remove the log file after test
+            std::fs::remove_file(&log_file).ok();
         })
     }
 }
diff --git a/ui/desktop/src/goosed.ts b/ui/desktop/src/goosed.ts
index 4b415891d..e5fb53f6a 100644
--- a/ui/desktop/src/goosed.ts
+++ b/ui/desktop/src/goosed.ts
@@ -3,6 +3,8 @@ import { createServer } from 'net';
 import os from 'node:os';
 import { getBinaryPath } from './utils/binaryPath';
 import log from './utils/logger';
+import { ChildProcessByStdio } from 'node:child_process';
+import { Readable } from 'node:stream';
 
 // Find an available port to start goosed on
 export const findAvailablePort = (): Promise<number> => {
@@ -19,20 +21,6 @@ export const findAvailablePort = (): Promise<number> => {
   });
 };
 
-// Function to fetch agent version from the server
-const fetchAgentVersion = async (port: number): Promise<string> => {
-  try {
-    const response = await fetch(`http://127.0.0.1:${port}/agent/versions`);
-    if (!response.ok) {
-      throw new Error(`HTTP error! status: ${response.status}`);
-    }
-    const data = await response.json();
-    return data.current_version;
-  } catch (error) {
-    log.error('Failed to fetch agent version:', error);
-    return 'unknown';
-  }
-};
 
 // Goose process manager. Take in the app, port, and directory to start goosed in.
 // Check if goosed server is ready by polling the status endpoint
@@ -58,7 +46,7 @@ const checkServerStatus = async (port: number, maxAttempts: number = 60, interva
   return false;
 };
 
-export const startGoosed = async (app, dir=null, env={}): Promise<[number, string, string]> => {
+export const startGoosed = async (app, dir=null, env={}): Promise<[number, string, ChildProcessByStdio<null, Readable, Readable>]> => {
   // we default to running goosed in home dir - if not specified
   const homeDir = os.homedir();
   if (!dir) {
@@ -128,10 +116,7 @@ export const startGoosed = async (app, dir=null, env={}): Promise<[number, strin
     goosedProcess.kill();
   });
 
-  // Wait for the server to start and fetch the agent version
-  await new Promise(resolve => setTimeout(resolve, 1000)); // Give the server time to start
-  const agentVersion = await fetchAgentVersion(port);
 
   log.info(`Goosed server successfully started on port ${port}`);
-  return [port, dir, agentVersion];
+  return [port, dir, goosedProcess];
 };
diff --git a/ui/desktop/src/main.ts b/ui/desktop/src/main.ts
index 8c77807fb..9c8434548 100644
--- a/ui/desktop/src/main.ts
+++ b/ui/desktop/src/main.ts
@@ -88,7 +88,6 @@ let appConfig = {
   GOOSE_API_HOST: 'http://127.0.0.1',
   GOOSE_PORT: 0,
   GOOSE_WORKING_DIR: '',
-  GOOSE_AGENT_VERSION: '',
   secretKey: generateSecretKey(),
 };
 
@@ -143,7 +142,7 @@ const createChat = async (app, query?: string, dir?: string, version?: string, d
   // Apply current environment settings before creating chat
   updateEnvironmentVariables(envToggles);
 
-  const [port, working_dir, agentVersion] = await startGoosed(app, dir);
+  const [port, working_dir, goosedProcess] = await startGoosed(app, dir);
 
   const mainWindow = new BrowserWindow({
     titleBarStyle: 'hidden',
@@ -162,7 +161,6 @@ const createChat = async (app, query?: string, dir?: string, version?: string, d
         ...appConfig, 
         GOOSE_PORT: port,
         GOOSE_WORKING_DIR: working_dir,
-        GOOSE_AGENT_VERSION: agentVersion,
         REQUEST_DIR: dir,
         DEEP_LINK: deepLink,
       })],
@@ -236,7 +234,10 @@ const createChat = async (app, query?: string, dir?: string, version?: string, d
   mainWindow.on('closed', () => {
     windowMap.delete(windowId);
     unregisterDevToolsShortcut();
+    goosedProcess.kill();
   });
+
+  
 };
 
 const createTray = () => {