From 9f7d6d428bff53af2703ca8f76fc596f600eb9ca Mon Sep 17 00:00:00 2001
From: Saoud Rizwan <7799382+saoudrizwan@users.noreply.github.com>
Date: Sat, 10 Aug 2024 03:41:47 -0400
Subject: [PATCH] Gracefully terminate running subprocess if user cancels task
 in the middle of a command

---
 src/ClaudeDev.ts                   | 14 +++++++++++++-
 src/providers/ClaudeDevProvider.ts |  8 ++------
 2 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/src/ClaudeDev.ts b/src/ClaudeDev.ts
index cf09fc5..1b61cba 100644
--- a/src/ClaudeDev.ts
+++ b/src/ClaudeDev.ts
@@ -242,8 +242,9 @@ export class ClaudeDev {
 	private askResponseText?: string
 	private askResponseImages?: string[]
 	private lastMessageTs?: number
+	private executeCommandRunningProcess?: ResultPromise
 	private providerRef: WeakRef<ClaudeDevProvider>
-	abort: boolean = false
+	private abort: boolean = false
 
 	constructor(
 		provider: ClaudeDevProvider,
@@ -382,6 +383,14 @@ export class ClaudeDev {
 		}
 	}
 
+	abortTask() {
+		this.abort = true // will stop any autonomously running promises
+		const runningProcessId = this.executeCommandRunningProcess?.pid
+		if (runningProcessId) {
+			treeKill(runningProcessId, "SIGTERM")
+		}
+	}
+
 	async executeTool(toolName: ToolName, toolInput: any, isLastWriteToFile: boolean = false): Promise<ToolResponse> {
 		switch (toolName) {
 			case "write_to_file":
@@ -821,6 +830,7 @@ export class ClaudeDev {
 			// also worth noting that execa`input` and the execa(command) have nuanced differences like the template literal version handles escaping for you, while with the function call, you need to be more careful about how arguments are passed, especially when using shell: true.
 			// execa returns a promise-like object that is both a promise and a Subprocess that has properties like stdin
 			const subprocess = execa({ shell: true, cwd: cwd })`${command}`
+			this.executeCommandRunningProcess = subprocess
 
 			subprocess.stdout?.on("data", (data) => {
 				if (data) {
@@ -854,6 +864,7 @@ export class ClaudeDev {
 			// the correct order of messages (although the webview is smart about
 			// grouping command_output messages despite any gaps anyways)
 			await delay(100)
+			this.executeCommandRunningProcess = undefined
 			// for attemptCompletion, we don't want to return the command output
 			if (returnEmptyStringOnSuccess) {
 				return ""
@@ -864,6 +875,7 @@ export class ClaudeDev {
 			let errorMessage = error.message || JSON.stringify(serializeError(error), null, 2)
 			const errorString = `Error executing command:\n${errorMessage}`
 			this.say("error", `Error executing command:\n${errorMessage}`) // TODO: in webview show code block for command errors
+			this.executeCommandRunningProcess = undefined
 			return errorString
 		}
 	}
diff --git a/src/providers/ClaudeDevProvider.ts b/src/providers/ClaudeDevProvider.ts
index 60dac29..536491e 100644
--- a/src/providers/ClaudeDevProvider.ts
+++ b/src/providers/ClaudeDevProvider.ts
@@ -332,12 +332,8 @@ export class ClaudeDevProvider implements vscode.WebviewViewProvider {
 	}
 
 	async clearTask() {
-		if (this.claudeDev) {
-			this.claudeDev.abort = true // will stop any agentically running promises
-			this.claudeDev = undefined // removes reference to it, so once promises end it will be garbage collected
-		}
-		// this.setApiConversationHistory(undefined)
-		// this.setClaudeMessages(undefined)
+		this.claudeDev?.abortTask()
+		this.claudeDev = undefined // removes reference to it, so once promises end it will be garbage collected
 	}
 
 	// Caching mechanism to keep track of webview messages + API conversation history per provider instance