From 7b9be980b938450d597e0f0c7705c6a4f080d109 Mon Sep 17 00:00:00 2001
From: Huakun Shen <huakun.shen@huakunshen.com>
Date: Tue, 12 Nov 2024 14:40:27 -0500
Subject: [PATCH] feat: implement extension spawned process auto cleanup

If extension doesn't kill the processes it spawns, Kunkun will auto clean up all spawned processes on exit/window close
---
 apps/desktop/src/lib/cmds/ext.ts               |  2 --
 apps/desktop/src/lib/stores/winExtMap.ts       | 18 ++++++++++++------
 apps/desktop/src/lib/utils/style.ts            |  1 -
 apps/desktop/src/routes/+layout.svelte         | 15 ++++++++++++++-
 apps/desktop/src/routes/+page.svelte           | 12 ++++--------
 .../routes/extension/ui-iframe/+page.svelte    |  1 -
 .../routes/extension/ui-worker/+page.svelte    |  6 +-----
 .../ui/src/components/main/ExtCmdsGroup.svelte |  6 +++---
 8 files changed, 34 insertions(+), 27 deletions(-)

diff --git a/apps/desktop/src/lib/cmds/ext.ts b/apps/desktop/src/lib/cmds/ext.ts
index b29abe9..5a6617a 100644
--- a/apps/desktop/src/lib/cmds/ext.ts
+++ b/apps/desktop/src/lib/cmds/ext.ts
@@ -63,7 +63,6 @@ export async function onCustomUiCmdSelect(
 			const newUrl = `http://${addr}`
 			url2 = `/extension/ui-iframe?url=${encodeURIComponent(newUrl)}&extPath=${encodeURIComponent(ext.extPath)}`
 		}
-		console.log("URL 2", url2)
 		const window = launchNewExtWindow(winLabel, url2, cmd.window)
 		window.onCloseRequested(async (event) => {
 			await winExtMap.unregisterExtensionFromWindow(winLabel)
@@ -81,7 +80,6 @@ export async function onCustomUiCmdSelect(
 			const newUrl = `http://${addr}`
 			url2 = `/extension/ui-iframe?url=${encodeURIComponent(newUrl)}&extPath=${encodeURIComponent(ext.extPath)}`
 		}
-		console.log("URL 2", url2)
 		goto(url2)
 	}
 	appState.clearSearchTerm()
diff --git a/apps/desktop/src/lib/stores/winExtMap.ts b/apps/desktop/src/lib/stores/winExtMap.ts
index 2bea26a..dfd3178 100644
--- a/apps/desktop/src/lib/stores/winExtMap.ts
+++ b/apps/desktop/src/lib/stores/winExtMap.ts
@@ -30,6 +30,7 @@ type API = {
 		dist?: string
 	}) => Promise<string>
 	unregisterExtensionFromWindow: (windowLabel: string) => Promise<void>
+	cleanupProcessesFromWindow: (windowLabel: string) => Promise<void>
 	registerProcess: (windowLabel: string, pid: number) => Promise<void>
 	unregisterProcess: (pid: number) => Promise<void>
 }
@@ -84,11 +85,10 @@ function createWinExtMapStore(): Writable<WinExtMap> & API {
 			if (winExtMap[windowLabel]) {
 				// clean up processes spawned by extension but not killed by itself
 				const extLabelMap = await getExtLabelMap() // realtime data from core process
-				Object.entries(extLabelMap).forEach(([label, ext]) => {
-					if (label === windowLabel) {
-						killProcesses(ext.processes)
-					}
-				})
+				if (extLabelMap[windowLabel]) {
+					console.log("kill processes", extLabelMap[windowLabel].processes)
+					killProcesses(extLabelMap[windowLabel].processes)
+				}
 				await unregisterExtensionWindow(windowLabel)
 				delete winExtMap[windowLabel]
 				store.set(winExtMap)
@@ -96,9 +96,15 @@ function createWinExtMapStore(): Writable<WinExtMap> & API {
 				warn(`Window ${windowLabel} does not have an extension registered`)
 			}
 		},
+		cleanupProcessesFromWindow: async (windowLabel: string) => {
+			const winExtMap = get(store)
+			if (winExtMap[windowLabel]) {
+				await killProcesses(winExtMap[windowLabel].pids)
+			}
+		},
 		registerProcess: async (windowLabel: string, pid: number) => {
 			const winExtMap = get(store)
-			registerExtensionSpawnedProcess(windowLabel, pid)
+			await registerExtensionSpawnedProcess(windowLabel, pid)
 			if (!winExtMap[windowLabel]) {
 				throw new Error(`Window ${windowLabel} does not have an extension registered`)
 			}
diff --git a/apps/desktop/src/lib/utils/style.ts b/apps/desktop/src/lib/utils/style.ts
index 12e3b17..f84a872 100644
--- a/apps/desktop/src/lib/utils/style.ts
+++ b/apps/desktop/src/lib/utils/style.ts
@@ -29,7 +29,6 @@ export function positionToTailwindClasses(position: Position) {
 			if (parseOutput.output.left) {
 				className += ` left-[${parseOutput.output.left / 4}rem]`
 			}
-			console.log(position, className)
 			return className
 	}
 }
diff --git a/apps/desktop/src/routes/+layout.svelte b/apps/desktop/src/routes/+layout.svelte
index 079021f..164a602 100644
--- a/apps/desktop/src/routes/+layout.svelte
+++ b/apps/desktop/src/routes/+layout.svelte
@@ -1,12 +1,13 @@
 <script lang="ts">
 	import AppContext from "@/components/context/AppContext.svelte"
 	import "../app.css"
-	import { appConfig, appState, extensions, quickLinks } from "@/stores"
+	import { appConfig, appState, extensions, quickLinks, winExtMap } from "@/stores"
 	import { initDeeplink } from "@/utils/deeplink"
 	import { updateAppHotkey } from "@/utils/hotkey"
 	import { globalKeyDownHandler, goBackOrCloseOnEscape } from "@/utils/key"
 	import { listenToWindowBlur } from "@/utils/tauri-events"
 	import { isInMainWindow } from "@/utils/window"
+	import { listenToKillProcessEvent, listenToRecordExtensionProcessEvent } from "@kksh/api/events"
 	import {
 		ModeWatcher,
 		themeConfigStore,
@@ -78,6 +79,18 @@
 				})
 			)
 			extensions.init()
+			unlisteners.push(
+				await listenToRecordExtensionProcessEvent(async (event) => {
+					console.log("record extension process event", event)
+					winExtMap.registerProcess(event.payload.windowLabel, event.payload.pid)
+				})
+			)
+			unlisteners.push(
+				await listenToKillProcessEvent((event) => {
+					console.log("kill process event", event)
+					winExtMap.unregisterProcess(event.payload.pid)
+				})
+			)
 		} else {
 		}
 		getCurrentWebviewWindow().show()
diff --git a/apps/desktop/src/routes/+page.svelte b/apps/desktop/src/routes/+page.svelte
index 10539a2..f23e668 100644
--- a/apps/desktop/src/routes/+page.svelte
+++ b/apps/desktop/src/routes/+page.svelte
@@ -5,13 +5,10 @@
 	import { systemCommands } from "@/cmds/system"
 	import { appConfig, appState, devStoreExts, installedStoreExts, quickLinks } from "@/stores"
 	import { cmdQueries } from "@/stores/cmdQuery"
-	import { getActiveElementNodeName, isKeyboardEventFromInputElement } from "@/utils/dom"
+	import { isKeyboardEventFromInputElement } from "@/utils/dom"
 	import Icon from "@iconify/svelte"
-	import { openDevTools, toggleDevTools } from "@kksh/api/commands"
-	import type { ExtPackageJsonExtra } from "@kksh/api/models"
-	import { isExtPathInDev } from "@kksh/extension/utils"
+	import { toggleDevTools } from "@kksh/api/commands"
 	import { Button, Command, DropdownMenu } from "@kksh/svelte5"
-	import type { AppConfig, AppState } from "@kksh/types"
 	import {
 		BuiltinCmds,
 		CustomCommandInput,
@@ -20,12 +17,11 @@
 		QuickLinks,
 		SystemCmds
 	} from "@kksh/ui/main"
-	import type { BuiltinCmd, CmdValue, CommandLaunchers } from "@kksh/ui/types"
+	import type { CmdValue } from "@kksh/ui/types"
 	import { cn, commandScore } from "@kksh/ui/utils"
 	import { getCurrentWebviewWindow } from "@tauri-apps/api/webviewWindow"
 	import { exit } from "@tauri-apps/plugin-process"
 	import { ArrowBigUpIcon, CircleXIcon, EllipsisVerticalIcon, RefreshCcwIcon } from "lucide-svelte"
-	import { onMount } from "svelte"
 
 	let inputEle: HTMLInputElement | null = null
 	function onKeyDown(event: KeyboardEvent) {
@@ -136,7 +132,7 @@
 						>
 							<Icon
 								icon={$appConfig.hmr ? "fontisto:toggle-on" : "fontisto:toggle-off"}
-								class={cn("mr-1 h-5 w-5", $appConfig.hmr ? "text-green-500" : "")}
+								class={cn("mr-1 h-5 w-5", { "text-green-500": $appConfig.hmr })}
 							/>
 							Toggle Dev Extension HMR
 						</DropdownMenu.Item>
diff --git a/apps/desktop/src/routes/extension/ui-iframe/+page.svelte b/apps/desktop/src/routes/extension/ui-iframe/+page.svelte
index b489a2b..654c840 100644
--- a/apps/desktop/src/routes/extension/ui-iframe/+page.svelte
+++ b/apps/desktop/src/routes/extension/ui-iframe/+page.svelte
@@ -112,7 +112,6 @@
 	} satisfies IApp
 
 	function onBackBtnClicked() {
-		console.log("onBackBtnClicked")
 		if (isInMainWindow()) {
 			goHome()
 		} else {
diff --git a/apps/desktop/src/routes/extension/ui-worker/+page.svelte b/apps/desktop/src/routes/extension/ui-worker/+page.svelte
index 4969a1d..50fb345 100644
--- a/apps/desktop/src/routes/extension/ui-worker/+page.svelte
+++ b/apps/desktop/src/routes/extension/ui-worker/+page.svelte
@@ -56,14 +56,9 @@
 	let loaded = $state(false)
 
 	async function goBack() {
-		console.log("goBack")
 		if (isInMainWindow()) {
-			console.log("goBack in main window")
-			// if in main window, then winExtMap store must contain this
-			// winExtMap.unregisterExtensionFromWindow(appWin.label)
 			goto("/")
 		} else {
-			console.log("goBack in webview window")
 			appWin.close()
 		}
 	}
@@ -235,6 +230,7 @@
 
 	onDestroy(() => {
 		unlistenRefreshWorkerExt?.()
+		winExtMap.unregisterExtensionFromWindow(appWin.label)
 		extensionLoadingBar = false
 		appState.setActionPanel(undefined)
 	})
diff --git a/packages/ui/src/components/main/ExtCmdsGroup.svelte b/packages/ui/src/components/main/ExtCmdsGroup.svelte
index 62102c8..df665b9 100644
--- a/packages/ui/src/components/main/ExtCmdsGroup.svelte
+++ b/packages/ui/src/components/main/ExtCmdsGroup.svelte
@@ -38,12 +38,12 @@
 			<IconMultiplexer icon={cmd.icon ?? ext.kunkun.icon} class="!h-5 !w-5 shrink-0" />
 			<span>{cmd.name}</span>
 		</span>
-		<span class="flex gap-2">
+		<span class="flex gap-1">
 			{#if isDev}
-				<Badge class="rounded-sm bg-green-500 px-1">Dev</Badge>
+				<Badge class="scale-75 rounded-sm bg-green-500 px-1">Dev</Badge>
 			{/if}
 			{#if hmr}
-				<Badge class="rounded-sm px-1">HMR</Badge>
+				<Badge class="scale-75 rounded-sm px-1">HMR</Badge>
 			{/if}
 		</span>
 	</Command.Item>