nes: feat: support next cursor line prediction

src/extension/xtab/common/promptCrafting.ts

@@ -16,6 +16,7 @@ import { illegalArgument } from '../../../util/vs/base/common/errors';
 import { Schemas } from '../../../util/vs/base/common/network';
 import { OffsetRange } from '../../../util/vs/editor/common/core/ranges/offsetRange';
 import { StringText } from '../../../util/vs/editor/common/core/text/abstractText';
+import { CurrentDocument } from './xtabCurrentDocument';
 
 export namespace PromptTags {
 	export const CURSOR = "<|cursor|>";
@@ -147,6 +148,9 @@ export const xtab275SystemPrompt = `Predict the next code edit based on user con
 
 export class PromptPieces {
 	constructor(
+		public readonly currentDocument: CurrentDocument,
+		public readonly editWindowLinesRange: OffsetRange,
+		public readonly areaAroundEditWindowLinesRange: OffsetRange,
 		public readonly activeDoc: StatelessNextEditDocument,
 		public readonly xtabHistory: readonly IXtabHistoryEntry[],
 		public readonly currentFileContent: string,
@@ -172,7 +176,7 @@ export function getUserPrompt(promptPieces: PromptPieces): string {
 
 	const currentFilePath = toUniquePath(activeDoc.id, activeDoc.workspaceRoot?.path);
 
-	const postScript = getPostScript(opts.promptingStrategy, currentFilePath);
+	const postScript = promptPieces.opts.includePostScript ? getPostScript(opts.promptingStrategy, currentFilePath) : '';
 
 	const mainPrompt = `${PromptTags.RECENT_FILES.start}
 ${recentlyViewedCodeSnippets}

src/extension/xtab/node/xtabProvider.ts

@@ -3,6 +3,7 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
+import { RequestType } from '@vscode/copilot-api';
 import { Raw } from '@vscode/prompt-tsx';
 import { ChatCompletionContentPartKind } from '@vscode/prompt-tsx/dist/base/output/rawTypes';
 import { FetchStreamSource } from '../../../platform/chat/common/chatMLFetcher';
@@ -34,6 +35,7 @@ import { IWorkspaceService } from '../../../platform/workspace/common/workspaceS
 import { raceFilter } from '../../../util/common/async';
 import * as errors from '../../../util/common/errors';
 import { Result } from '../../../util/common/result';
+import { TokenizerType } from '../../../util/common/tokenizer';
 import { createTracer, ITracer } from '../../../util/common/tracing';
 import { AsyncIterableObject, DeferredPromise, raceTimeout, timeout } from '../../../util/vs/base/common/async';
 import { CancellationToken } from '../../../util/vs/base/common/cancellation';
@@ -287,6 +289,9 @@ export class XtabProvider implements IStatelessNextEditProvider {
 		}
 
 		const promptPieces = new PromptPieces(
+			currentDocument,
+			editWindowLinesRange,
+			areaAroundEditWindowLinesRange,
 			activeDocument,
 			request.xtabEditHistory,
 			taggedCurrentFileContent,
@@ -332,6 +337,7 @@ export class XtabProvider implements IStatelessNextEditProvider {
 			cursorOriginalLinesOffset,
 			cursorLineOffset,
 			editWindowLinesRange,
+			promptPieces,
 			prediction,
 			{
 				shouldRemoveCursorTagFromResponse,
@@ -351,7 +357,7 @@ export class XtabProvider implements IStatelessNextEditProvider {
 		currentDocument: CurrentDocument,
 		editWindowLinesRange: OffsetRange,
 		areaAroundEditWindowLinesRange: OffsetRange,
-		promptOptions: ModelConfig,
+		promptOptions: xtabPromptOptions.PromptOptions,
 		computeTokens: (s: string) => number,
 		opts: {
 			includeLineNumbers: boolean;
@@ -533,6 +539,7 @@ export class XtabProvider implements IStatelessNextEditProvider {
 		cursorOriginalLinesOffset: number,
 		cursorLineOffset: number, // cursor offset within the line it's in; 1-based
 		editWindowLineRange: OffsetRange,
+		promptPieces: PromptPieces,
 		prediction: Prediction | undefined,
 		opts: {
 			promptingStrategy: xtabPromptOptions.PromptingStrategy | undefined;
@@ -679,7 +686,7 @@ export class XtabProvider implements IStatelessNextEditProvider {
 			const trimmedLines = firstLine.value.trim();
 
 			if (trimmedLines === ResponseTags.NO_CHANGE.start) {
-				this.pushNoSuggestionsOrRetry(request, editWindow, pushEdit, delaySession, logContext, cancellationToken, telemetryBuilder, opts.retryState);
+				await this.pushNoSuggestionsOrRetry(request, editWindow, promptPieces, pushEdit, delaySession, logContext, cancellationToken, telemetryBuilder, opts.retryState);
 				return;
 			}
 
@@ -807,7 +814,7 @@ export class XtabProvider implements IStatelessNextEditProvider {
 				if (hadEdits) {
 					pushEdit(Result.error(new NoNextEditReason.NoSuggestions(request.documentBeforeEdits, editWindow)));
 				} else {
-					this.pushNoSuggestionsOrRetry(request, editWindow, pushEdit, delaySession, logContext, cancellationToken, telemetryBuilder, opts.retryState);
+					await this.pushNoSuggestionsOrRetry(request, editWindow, promptPieces, pushEdit, delaySession, logContext, cancellationToken, telemetryBuilder, opts.retryState);
 				}
 
 			} catch (err) {
@@ -818,9 +825,10 @@ export class XtabProvider implements IStatelessNextEditProvider {
 		})();
 	}
 
-	private pushNoSuggestionsOrRetry(
+	private async pushNoSuggestionsOrRetry(
 		request: StatelessNextEditRequest,
 		editWindow: OffsetRange,
+		promptPieces: PromptPieces,
 		pushEdit: PushEdit,
 		delaySession: DelaySession,
 		logContext: InlineEditRequestLogContext,
@@ -836,6 +844,17 @@ export class XtabProvider implements IStatelessNextEditProvider {
 			return;
 		}
 
+		// FIXME@ulugbekna: think out how it works with retrying logic
+		if (this.configService.getExperimentBasedConfig(ConfigKey.Internal.InlineEditsNextCursorPredictionEnabled, this.expService)) {
+			// FIXME@ulugbekna: possibly convert from 1-based to 0-based
+			const nextCursorLine = await this.predictNextCursorPosition(promptPieces);
+			if (nextCursorLine) {
+				this.tracer.trace(`Predicted next cursor line: ${nextCursorLine}`);
+				this.doGetNextEditWithSelection(request, new Range(nextCursorLine, 1, nextCursorLine, 1), pushEdit, delaySession, logContext, cancellationToken, telemetryBuilder, RetryState.NotRetrying);
+				return;
+			}
+		}
+
 		pushEdit(Result.error(new NoNextEditReason.NoSuggestions(request.documentBeforeEdits, editWindow)));
 		return;
 	}
@@ -977,7 +996,8 @@ export class XtabProvider implements IStatelessNextEditProvider {
 				maxTokens: this.configService.getExperimentBasedConfig(ConfigKey.Internal.InlineEditsXtabDiffMaxTokens, this.expService),
 				onlyForDocsInPrompt: this.configService.getExperimentBasedConfig(ConfigKey.Internal.InlineEditsXtabDiffOnlyForDocsInPrompt, this.expService),
 				useRelativePaths: this.configService.getExperimentBasedConfig(ConfigKey.Internal.InlineEditsXtabDiffUseRelativePaths, this.expService),
-			}
+			},
+			includePostScript: true,
 		};
 
 		const overridingModelConfig = this.configService.getConfig(ConfigKey.Internal.InlineEditsXtabProviderModelConfiguration);
@@ -996,6 +1016,120 @@ export class XtabProvider implements IStatelessNextEditProvider {
 		return sourcedModelConfig;
 	}
 
+	private async predictNextCursorPosition(promptPieces: PromptPieces) {
+
+		const tracer = this.tracer.sub('predictNextCursorPosition');
+
+		const systemMessage = 'Your task is to predict the next line number in the current file where the developer is most likely to make their next edit, using the provided context.';
+
+		const currentFileContentR = this.constructTaggedFile(
+			promptPieces.currentDocument,
+			promptPieces.editWindowLinesRange,
+			promptPieces.areaAroundEditWindowLinesRange,
+			promptPieces.opts,
+			XtabProvider.computeTokens,
+			{ includeLineNumbers: true }
+		);
+
+		if (currentFileContentR.isError()) {
+			tracer.trace(`Failed to construct tagged file: ${currentFileContentR.err}`);
+			return;
+		}
+
+		const { taggedCurrentFileR: { taggedCurrentFileContent }, areaAroundCodeToEdit } = currentFileContentR.val;
+
+		const newPromptPieces = new PromptPieces(
+			promptPieces.currentDocument,
+			promptPieces.editWindowLinesRange,
+			promptPieces.areaAroundEditWindowLinesRange,
+			promptPieces.activeDoc,
+			promptPieces.xtabHistory,
+			taggedCurrentFileContent,
+			areaAroundCodeToEdit,
+			promptPieces.langCtx,
+			XtabProvider.computeTokens,
+			{
+				...promptPieces.opts,
+				includePostScript: false,
+			}
+		);
+
+		const userMessage = getUserPrompt(newPromptPieces);
+
+		const messages = constructMessages({
+			systemMsg: systemMessage,
+			userMsg: userMessage
+		});
+
+		const modelName = this.configService.getExperimentBasedConfig(ConfigKey.Internal.InlineEditsNextCursorPredictionModelName, this.expService);
+		if (modelName === undefined) {
+			tracer.trace('Model name for cursor prediction is not defined; skipping prediction');
+			return;
+		}
+
+		const url = this.configService.getConfig(ConfigKey.Internal.InlineEditsNextCursorPredictionUrl);
+		const secretKey = this.configService.getConfig(ConfigKey.Internal.InlineEditsNextCursorPredictionApiKey);
+
+		const endpoint = this.instaService.createInstance(ChatEndpoint, {
+			id: modelName,
+			name: 'nes.nextCursorPosition',
+			urlOrRequestMetadata: url ? url : { type: RequestType.ProxyChatCompletions },
+			model_picker_enabled: false,
+			is_chat_default: false,
+			is_chat_fallback: false,
+			version: '',
+			capabilities: {
+				type: 'chat',
+				family: '',
+				tokenizer: TokenizerType.CL100K,
+				limits: undefined,
+				supports: {
+					parallel_tool_calls: false,
+					tool_calls: false,
+					streaming: true,
+					vision: false,
+					prediction: false,
+					thinking: false
+				}
+			},
+		});
+
+		const response = await endpoint.makeChatRequest2(
+			{
+				messages,
+				debugName: 'nes.nextCursorPosition',
+				finishedCb: undefined,
+				location: ChatLocation.Other,
+				requestOptions: secretKey ? {
+					secretKey,
+				} : undefined,
+			},
+			CancellationToken.None
+		);
+
+		if (response.type !== ChatFetchResponseType.Success) {
+			return;
+		}
+
+		try {
+			const trimmed = response.value.trim();
+			const lineNumber = parseInt(trimmed, 10);
+			if (isNaN(lineNumber) || lineNumber < 0) {
+				throw new Error(`parsed line number is NaN or negative: ${trimmed}`);
+			}
+			if (lineNumber > promptPieces.currentDocument.lines.length) {
+				this.tracer.trace(`Predicted line number ${lineNumber} is out of bounds, document has ${promptPieces.currentDocument.lines.length} lines`);
+				return undefined;
+			}
+
+			return lineNumber;
+		} catch (err) {
+			tracer.trace(`Failed to parse predicted line number from response '${response.value}': ${err}`);
+			return undefined;
+		}
+
+	}
+
 	private determinePromptingStrategy(): xtabPromptOptions.PromptingStrategy | undefined {
 		const isXtabUnifiedModel = this.configService.getExperimentBasedConfig(ConfigKey.Internal.InlineEditsXtabUseUnifiedModel, this.expService);
 		const isCodexV21NesUnified = this.configService.getExperimentBasedConfig(ConfigKey.Internal.InlineEditsXtabCodexV21NesUnified, this.expService);

src/platform/configuration/common/configurationService.ts

@@ -664,6 +664,10 @@ export namespace ConfigKey {
 		export const InlineEditsHideInternalInterface = defineValidatedSetting<boolean>('chat.advanced.inlineEdits.hideInternalInterface', vBoolean(), false, INTERNAL_RESTRICTED);
 		export const InlineEditsLogCancelledRequests = defineValidatedSetting<boolean>('chat.advanced.inlineEdits.logCancelledRequests', vBoolean(), false, INTERNAL_RESTRICTED);
 		export const InlineEditsUnification = defineExpSetting<boolean>('chat.advanced.inlineEdits.unification', false, INTERNAL_RESTRICTED);
+		export const InlineEditsNextCursorPredictionEnabled = defineExpSetting<boolean>('chat.advanced.inlineEdits.nextCursorPrediction.enabled', false, INTERNAL_RESTRICTED);
+		export const InlineEditsNextCursorPredictionModelName = defineExpSetting<string | undefined>('chat.advanced.inlineEdits.nextCursorPrediction.modelName', undefined, INTERNAL_RESTRICTED);
+		export const InlineEditsNextCursorPredictionUrl = defineValidatedSetting<string | undefined>('chat.advanced.inlineEdits.nextCursorPrediction.url', vString(), undefined, INTERNAL_RESTRICTED);
+		export const InlineEditsNextCursorPredictionApiKey = defineValidatedSetting<string | undefined>('chat.advanced.inlineEdits.nextCursorPrediction.apiKey', vString(), undefined, INTERNAL_RESTRICTED);
 		export const InlineEditsXtabProviderUrl = defineValidatedSetting<string | undefined>('chat.advanced.inlineEdits.xtabProvider.url', vString(), undefined, INTERNAL_RESTRICTED);
 		export const InlineEditsXtabProviderApiKey = defineValidatedSetting<string | undefined>('chat.advanced.inlineEdits.xtabProvider.apiKey', vString(), undefined, INTERNAL_RESTRICTED);
 		export const InlineEditsXtabProviderModelConfiguration = defineValidatedSetting<xtabPromptOptions.ModelConfiguration | undefined>('chat.advanced.inlineEdits.xtabProvider.modelConfiguration', xtabPromptOptions.MODEL_CONFIGURATION_VALIDATOR, undefined, INTERNAL_RESTRICTED);

src/platform/inlineEdits/common/dataTypes/xtabPromptOptions.ts

@@ -40,6 +40,7 @@ export type PromptOptions = {
 	readonly recentlyViewedDocuments: RecentlyViewedDocumentsOptions;
 	readonly languageContext: LanguageContextOptions;
 	readonly diffHistory: DiffHistoryOptions;
+	readonly includePostScript: boolean;
 }
 
 /**
@@ -81,6 +82,7 @@ export const DEFAULT_OPTIONS: PromptOptions = {
 		onlyForDocsInPrompt: false,
 		useRelativePaths: false,
 	},
+	includePostScript: true,
 };
 
 // TODO: consider a better per language setting/experiment approach