From e740a5ccb06c7097c58c69781cf83e6e2d98ee16 Mon Sep 17 00:00:00 2001
From: Daniel Imms <2193314+Tyriar@users.noreply.github.com>
Date: Fri, 17 Nov 2023 09:30:56 -0800
Subject: [PATCH] Adjust command finished marker in started adjustment

Part of #196634
---
 .../commandDetectionCapability.ts             | 31 ++++++++++++++-----
 1 file changed, 24 insertions(+), 7 deletions(-)

diff --git a/src/vs/platform/terminal/common/capabilities/commandDetectionCapability.ts b/src/vs/platform/terminal/common/capabilities/commandDetectionCapability.ts
index d4d2892df167a..da16580b4b4e1 100644
--- a/src/vs/platform/terminal/common/capabilities/commandDetectionCapability.ts
+++ b/src/vs/platform/terminal/common/capabilities/commandDetectionCapability.ts
@@ -32,7 +32,6 @@ export class CommandDetectionCapability extends Disposable implements ICommandDe
 	private __isCommandStorageDisabled: boolean = false;
 	private _handleCommandStartOptions?: IHandleCommandOptions;
 
-
 	private _commitCommandFinished?: RunOnceScheduler;
 
 	private _ptyHeuristicsHooks: ICommandDetectionHeuristicsHooks;
@@ -99,6 +98,10 @@ export class CommandDetectionCapability extends Disposable implements ICommandDe
 			get commandMarkers() { return that._commandMarkers; }
 			set commandMarkers(value) { that._commandMarkers = value; }
 			get clearCommandsInViewport() { return that._clearCommandsInViewport.bind(that); }
+			commitCommandFinished() {
+				that._commitCommandFinished?.flush();
+				that._commitCommandFinished = undefined;
+			}
 		};
 		this._ptyHeuristics = this._register(new MandatoryMutableDisposable(new UnixPtyHeuristics(this._terminal, this, this._ptyHeuristicsHooks, this._logService)));
 
@@ -172,6 +175,10 @@ export class CommandDetectionCapability extends Disposable implements ICommandDe
 					get commandMarkers() { return that._commandMarkers; }
 					set commandMarkers(value) { that._commandMarkers = value; }
 					get clearCommandsInViewport() { return that._clearCommandsInViewport.bind(that); }
+					commitCommandFinished() {
+						that._commitCommandFinished?.flush();
+						that._commitCommandFinished = undefined;
+					}
 				},
 				this._logService
 			);
@@ -235,8 +242,6 @@ export class CommandDetectionCapability extends Disposable implements ICommandDe
 			this._logService.debug('CommandDetectionCapability#handlePromptStart adjusted commandFinished', `${lastCommand.endMarker.line} -> ${lastCommand.executedMarker.line + 1}`);
 			lastCommand.endMarker = cloneMarker(this._terminal, lastCommand.executedMarker, 1);
 		}
-		this._commitCommandFinished?.flush();
-		this._commitCommandFinished = undefined;
 
 		this._currentCommand.promptStartMarker = options?.marker || (lastCommand?.endMarker ? cloneMarker(this._terminal, lastCommand.endMarker) : this._terminal.registerMarker(0));
 
@@ -408,6 +413,7 @@ interface ICommandDetectionHeuristicsHooks {
 	commandMarkers: IMarker[];
 
 	clearCommandsInViewport(): void;
+	commitCommandFinished(): void;
 }
 
 type IPtyHeuristics = (
@@ -439,6 +445,8 @@ class UnixPtyHeuristics extends Disposable {
 	}
 
 	async handleCommandStart(options?: IHandleCommandOptions) {
+		this._hooks.commitCommandFinished();
+
 		const currentCommand = this._capability.currentCommand;
 		currentCommand.commandStartX = this._terminal.buffer.active.cursorX;
 		currentCommand.commandStartMarker = options?.marker || this._terminal.registerMarker(0);
@@ -507,6 +515,10 @@ class WindowsPtyHeuristics extends Disposable {
 
 	private _recentlyPerformedCsiJ = false;
 
+	private _tryAdjustCommandStartMarkerScheduler?: RunOnceScheduler;
+	private _tryAdjustCommandStartMarkerScannedLineCount: number = 0;
+	private _tryAdjustCommandStartMarkerPollCount: number = 0;
+
 	constructor(
 		private readonly _terminal: Terminal,
 		private readonly _capability: CommandDetectionCapability,
@@ -592,10 +604,6 @@ class WindowsPtyHeuristics extends Disposable {
 		}
 	}
 
-	private _tryAdjustCommandStartMarkerScheduler?: RunOnceScheduler;
-	private _tryAdjustCommandStartMarkerScannedLineCount: number = 0;
-	private _tryAdjustCommandStartMarkerPollCount: number = 0;
-
 	async handleCommandStart() {
 		this._capability.currentCommand.commandStartX = this._terminal.buffer.active.cursorX;
 
@@ -661,6 +669,13 @@ class WindowsPtyHeuristics extends Disposable {
 						this._logService.debug('CommandDetectionCapability#_tryAdjustCommandStartMarker adjusted promptStart', `${this._capability.currentCommand.promptStartMarker?.line} -> ${this._capability.currentCommand.commandStartMarker.line}`);
 						this._capability.currentCommand.promptStartMarker?.dispose();
 						this._capability.currentCommand.promptStartMarker = cloneMarker(this._terminal, this._capability.currentCommand.commandStartMarker);
+						// Adjust the last command if it's not in the same position as the following
+						// prompt start marker
+						const lastCommand = this._capability.commands.at(-1);
+						if (lastCommand && this._capability.currentCommand.commandStartMarker.line !== lastCommand.endMarker?.line) {
+							lastCommand.endMarker?.dispose();
+							lastCommand.endMarker = cloneMarker(this._terminal, this._capability.currentCommand.commandStartMarker);
+						}
 					}
 					// use the regex to set the position as it's possible input has occurred
 					this._capability.currentCommand.commandStartX = adjustedPrompt.length;
@@ -692,6 +707,8 @@ class WindowsPtyHeuristics extends Disposable {
 			this._tryAdjustCommandStartMarkerScheduler = undefined;
 		}
 
+		this._hooks.commitCommandFinished();
+
 		if (!this._capability.currentCommand.commandExecutedMarker) {
 			this._onCursorMoveListener.value = this._terminal.onCursorMove(() => {
 				if (this._hooks.commandMarkers.length === 0 || this._hooks.commandMarkers[this._hooks.commandMarkers.length - 1].line !== this._terminal.buffer.active.cursorY) {