persist command in pty service, refactor how commands are restored on window reload for inline chat terminal (#274417)

Author: meganrogge

src/vs/platform/terminal/common/capabilities/capabilities.ts

@@ -251,6 +251,12 @@ export interface ICommandDetectionCapability {
 	 * always be present when running the _builtin_ SI scripts.
 	 */
 	setCommandLine(commandLine: string, isTrusted: boolean): void;
+	/**
+	 * Sets the command ID to use for the next command that starts.
+	 * This allows pre-assigning an ID before the shell sends the command start sequence,
+	 * which is useful for linking commands across renderer and ptyHost.
+	 */
+	setNextCommandId(command: string, commandId: string): void;
 	serialize(): ISerializedCommandDetectionCapability;
 	deserialize(serialized: ISerializedCommandDetectionCapability): void;
 }
@@ -270,6 +276,12 @@ export interface IHandleCommandOptions {
 	 * Properties for the mark
 	 */
 	markProperties?: IMarkProperties;
+
+	/**
+	 * An optional predefined command ID. When provided, this ID will be used instead of
+	 * generating a new one, allowing commands to be linked across renderer and ptyHost.
+	 */
+	commandId?: string;
 }
 
 export interface INaiveCwdDetectionCapability {
@@ -291,7 +303,7 @@ interface IBaseTerminalCommand {
 	isTrusted: boolean;
 	timestamp: number;
 	duration: number;
-	id: string;
+	id: string | undefined;
 
 	// Optional serializable
 	cwd: string | undefined;

src/vs/platform/terminal/common/capabilities/commandDetection/terminalCommand.ts

@@ -6,15 +6,14 @@
 import { IMarkProperties, ISerializedTerminalCommand, ITerminalCommand } from '../capabilities.js';
 import { ITerminalOutputMatcher, ITerminalOutputMatch } from '../../terminal.js';
 import type { IBuffer, IBufferLine, IMarker, Terminal } from '@xterm/headless';
-import { generateUuid } from '../../../../../base/common/uuid.js';
 
 export interface ITerminalCommandProperties {
 	command: string;
 	commandLineConfidence: 'low' | 'medium' | 'high';
 	isTrusted: boolean;
 	timestamp: number;
 	duration: number;
-	id: string;
+	id: string | undefined;
 	marker: IMarker | undefined;
 	cwd: string | undefined;
 	exitCode: number | undefined;
@@ -276,7 +275,7 @@ export class PartialTerminalCommand implements ICurrentPartialCommand {
 	cwd?: string;
 	command?: string;
 	commandLineConfidence?: 'low' | 'medium' | 'high';
-	id: string;
+	id: string | undefined;
 
 	isTrusted?: boolean;
 	isInvalid?: boolean;
@@ -285,8 +284,7 @@ export class PartialTerminalCommand implements ICurrentPartialCommand {
 		private readonly _xterm: Terminal,
 		id?: string
 	) {
-		//TODO: this does not restore properly due to conflicting with the one created in the. PtyHost
-		this.id = id ?? generateUuid();
+		this.id = id;
 	}
 
 	serialize(cwd: string | undefined): ISerializedTerminalCommand | undefined {

src/vs/platform/terminal/common/capabilities/commandDetectionCapability.ts

@@ -35,6 +35,7 @@ export class CommandDetectionCapability extends Disposable implements ICommandDe
 	private _handleCommandStartOptions?: IHandleCommandOptions;
 	private _hasRichCommandDetection: boolean = false;
 	get hasRichCommandDetection() { return this._hasRichCommandDetection; }
+	private _nextCommandId: { command: string; commandId: string | undefined } | undefined;
 
 	private _ptyHeuristicsHooks: ICommandDetectionHeuristicsHooks;
 	private readonly _ptyHeuristics: MandatoryMutableDisposable<IPtyHeuristics>;
@@ -342,6 +343,14 @@ export class CommandDetectionCapability extends Disposable implements ICommandDe
 		this._ptyHeuristics.value.handleCommandStart(options);
 	}
 
+	/**
+	 * Sets the command ID to use for the next command that starts.
+	 * This is useful when you want to pre-assign an ID before the shell sends the command start sequence.
+	 */
+	setNextCommandId(command: string, commandId: string): void {
+		this._nextCommandId = { command, commandId };
+	}
+
 	handleCommandExecuted(options?: IHandleCommandOptions): void {
 		this._ptyHeuristics.value.handleCommandExecuted(options);
 		this._currentCommand.markExecutedTime();
@@ -355,7 +364,20 @@ export class CommandDetectionCapability extends Disposable implements ICommandDe
 		if (!this._currentCommand.commandExecutedMarker) {
 			this.handleCommandExecuted();
 		}
-
+		// If a custom command ID is provided, use it for the current command
+		// Otherwise, check if there's a pending next command ID
+		if (options?.commandId) {
+			this._currentCommand.id = options.commandId;
+			this._nextCommandId = undefined; // Clear the pending ID
+		} else if (
+			this._nextCommandId &&
+			typeof this.currentCommand.command === 'string' &&
+			typeof this._nextCommandId.command === 'string' &&
+			this.currentCommand.command.trim() === this._nextCommandId.command.trim()
+		) {
+			this._currentCommand.id = this._nextCommandId.commandId;
+			this._nextCommandId = undefined; // Clear after use
+		}
 		this._currentCommand.markFinishedTime();
 		this._ptyHeuristics.value.preHandleCommandFinished?.();
 
@@ -392,7 +414,9 @@ export class CommandDetectionCapability extends Disposable implements ICommandDe
 			this._logService.debug('CommandDetectionCapability#onCommandFinished', newCommand);
 			this._onCommandFinished.fire(newCommand);
 		}
-		this._currentCommand = new PartialTerminalCommand(this._terminal);
+		// Create new command for next execution, preserving command ID if one was specified
+		const nextCommandId = this._handleCommandStartOptions?.commandId;
+		this._currentCommand = new PartialTerminalCommand(this._terminal, nextCommandId);
 		this._handleCommandStartOptions = undefined;
 	}
 

src/vs/platform/terminal/common/terminal.ts

@@ -339,6 +339,7 @@ export interface IPtyService {
 	getInitialCwd(id: number): Promise<string>;
 	getCwd(id: number): Promise<string>;
 	acknowledgeDataEvent(id: number, charCount: number): Promise<void>;
+	setNextCommandId(id: number, commandLine: string, commandId: string): Promise<void>;
 	setUnicodeVersion(id: number, version: '6' | '11'): Promise<void>;
 	processBinary(id: number, data: string): Promise<void>;
 	/** Confirm the process is _not_ an orphan. */
@@ -816,6 +817,12 @@ export interface ITerminalChildProcess {
 	 */
 	acknowledgeDataEvent(charCount: number): void;
 
+	/**
+	 * Pre-assigns the command identifier that should be associated with the next command detected by
+	 * shell integration. This keeps the pty host and renderer command stores aligned.
+	 */
+	setNextCommandId(commandLine: string, commandId: string): Promise<void>;
+
 	/**
 	 * Sets the unicode version for the process, this drives the size of some characters in the
 	 * xterm-headless instance.
@@ -976,6 +983,8 @@ export interface IShellIntegration {
 	readonly onDidChangeSeenSequences: Event<ReadonlySet<string>>;
 
 	deserialize(serialized: ISerializedCommandDetectionCapability): void;
+
+	setNextCommandId(command: string, commandId: string): void;
 }
 
 export interface IDecorationAddon {

src/vs/platform/terminal/common/xterm/shellIntegrationAddon.ts

@@ -376,6 +376,12 @@ export class ShellIntegrationAddon extends Disposable implements IShellIntegrati
 		this._createOrGetBufferMarkDetection(terminal).getMark(vscodeMarkerId);
 	}
 
+	setNextCommandId(command: string, commandId: string): void {
+		if (this._terminal) {
+			this._createOrGetCommandDetection(this._terminal).setNextCommandId(command, commandId);
+		}
+	}
+
 	private _markSequenceSeen(sequence: string) {
 		if (!this._seenSequences.has(sequence)) {
 			this._seenSequences.add(sequence);

src/vs/platform/terminal/node/ptyHostService.ts

@@ -266,6 +266,9 @@ export class PtyHostService extends Disposable implements IPtyHostService {
 	setUnicodeVersion(id: number, version: '6' | '11'): Promise<void> {
 		return this._proxy.setUnicodeVersion(id, version);
 	}
+	setNextCommandId(id: number, commandLine: string, commandId: string): Promise<void> {
+		return this._proxy.setNextCommandId(id, commandLine, commandId);
+	}
 	getInitialCwd(id: number): Promise<string> {
 		return this._proxy.getInitialCwd(id);
 	}

src/vs/platform/terminal/node/ptyService.ts

@@ -455,6 +455,11 @@ export class PtyService extends Disposable implements IPtyService {
 	async setUnicodeVersion(id: number, version: '6' | '11'): Promise<void> {
 		return this._throwIfNoPty(id).setUnicodeVersion(version);
 	}
+
+	@traceRpc
+	async setNextCommandId(id: number, commandLine: string, commandId: string): Promise<void> {
+		return this._throwIfNoPty(id).setNextCommandId(commandLine, commandId);
+	}
 	@traceRpc
 	async getLatency(): Promise<IPtyHostLatencyMeasurement[]> {
 		return [];
@@ -892,6 +897,11 @@ class PersistentTerminalProcess extends Disposable {
 		this._serializer.setUnicodeVersion?.(version);
 		// TODO: Pass in unicode version in ctor
 	}
+
+	async setNextCommandId(commandLine: string, commandId: string): Promise<void> {
+		this._serializer.setNextCommandId?.(commandLine, commandId);
+	}
+
 	acknowledgeDataEvent(charCount: number): void {
 		if (this._inReplay) {
 			return;
@@ -1039,6 +1049,10 @@ class XtermSerializer implements ITerminalSerializer {
 		this._xterm.clear();
 	}
 
+	setNextCommandId(commandLine: string, commandId: string): void {
+		this._shellIntegrationAddon.setNextCommandId(commandLine, commandId);
+	}
+
 	async generateReplayEvent(normalBufferOnly?: boolean, restoreToLastReviveBuffer?: boolean): Promise<IPtyHostProcessReplayEvent> {
 		const serialize = new (await this._getSerializeConstructor());
 		this._xterm.loadAddon(serialize);
@@ -1126,4 +1140,5 @@ interface ITerminalSerializer {
 	clearBuffer(): void;
 	generateReplayEvent(normalBufferOnly?: boolean, restoreToLastReviveBuffer?: boolean): Promise<IPtyHostProcessReplayEvent>;
 	setUnicodeVersion?(version: '6' | '11'): void;
+	setNextCommandId?(commandLine: string, commandId: string): void;
 }

src/vs/platform/terminal/node/terminalProcess.ts

@@ -614,6 +614,10 @@ export class TerminalProcess extends Disposable implements ITerminalChildProcess
 		// No-op
 	}
 
+	async setNextCommandId(commandLine: string, commandId: string): Promise<void> {
+		// No-op: command IDs are tracked on the renderer and serializer only.
+	}
+
 	getInitialCwd(): Promise<string> {
 		return Promise.resolve(this._initialCwd);
 	}

src/vs/server/node/remoteTerminalChannel.ts

@@ -151,6 +151,7 @@ export class RemoteTerminalChannel extends Disposable implements IServerChannel<
 			case RemoteTerminalChannelRequest.ReviveTerminalProcesses: return this._ptyHostService.reviveTerminalProcesses.apply(this._ptyHostService, args);
 			case RemoteTerminalChannelRequest.GetRevivedPtyNewId: return this._ptyHostService.getRevivedPtyNewId.apply(this._ptyHostService, args);
 			case RemoteTerminalChannelRequest.SetUnicodeVersion: return this._ptyHostService.setUnicodeVersion.apply(this._ptyHostService, args);
+			case RemoteTerminalChannelRequest.SetNextCommandId: return this._ptyHostService.setNextCommandId.apply(this._ptyHostService, args);
 			case RemoteTerminalChannelRequest.ReduceConnectionGraceTime: return this._reduceConnectionGraceTime();
 			case RemoteTerminalChannelRequest.UpdateIcon: return this._ptyHostService.updateIcon.apply(this._ptyHostService, args);
 			case RemoteTerminalChannelRequest.UpdateTitle: return this._ptyHostService.updateTitle.apply(this._ptyHostService, args);

src/vs/workbench/api/common/extHostTerminalService.ts

@@ -363,6 +363,10 @@ class ExtHostPseudoterminal implements ITerminalChildProcess {
 		// No-op, xterm-headless isn't used for extension owned terminals.
 	}
 
+	async setNextCommandId(commandLine: string, commandId: string): Promise<void> {
+		// No-op, command IDs are only tracked on the renderer for extension terminals.
+	}
+
 	getInitialCwd(): Promise<string> {
 		return Promise.resolve('');
 	}