nes-datagen: add cursor-jump (NCLP) sample task

extensions/copilot/src/extension/xtab/node/xtabNextCursorPredictor.ts

@@ -3,6 +3,7 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
+import { Raw } from '@vscode/prompt-tsx';
 import { RequestType } from '@vscode/copilot-api';
 import { ChatFetchResponseType, ChatLocation } from '../../../platform/chat/common/commonTypes';
 import { ConfigKey, IConfigurationService } from '../../../platform/configuration/common/configurationService';
@@ -11,7 +12,6 @@ import { ChatEndpoint } from '../../../platform/endpoint/node/chatEndpoint';
 import { NextCursorLinePrediction } from '../../../platform/inlineEdits/common/dataTypes/nextCursorLinePrediction';
 import * as xtabPromptOptions from '../../../platform/inlineEdits/common/dataTypes/xtabPromptOptions';
 import { DEFAULT_CURSOR_PREDICTION_LINT_OPTIONS, parseLintOptionString } from '../../../platform/inlineEdits/common/dataTypes/xtabPromptOptions';
-import { StatelessNextEditTelemetryBuilder } from '../../../platform/inlineEdits/common/statelessNextEditProvider';
 import { ILanguageDiagnosticsService } from '../../../platform/languages/common/languageDiagnosticsService';
 import { ILogger } from '../../../platform/log/common/logService';
 import { OptionalChatRequestParams } from '../../../platform/networking/common/fetch';
@@ -28,6 +28,7 @@ import { OffsetRange } from '../../../util/vs/editor/common/core/ranges/offsetRa
 import { IInstantiationService } from '../../../util/vs/platform/instantiation/common/instantiation';
 import { LintErrors } from '../common/lintErrors';
 import { constructTaggedFile, getUserPrompt, PromptPieces } from '../common/promptCrafting';
+import type { RequestTracingContext } from './xtabProvider';
 import { constructMessages } from './xtabUtils';
 
 export type CursorJumpPrediction =
@@ -36,6 +37,14 @@ export type CursorJumpPrediction =
 
 const DEFAULT_CURSOR_JUMP_MODEL_NAME = 'copilot-suggestions-himalia-001';
 
+/**
+ * System prompt used for the cursor-jump (next-cursor-line) prediction model.
+ * Kept as an exported constant so that training-data generation can mirror it
+ * verbatim — drift between this string and the datagen prompt would corrupt
+ * the training distribution.
+ */
+export const NEXT_CURSOR_PREDICTION_SYSTEM_MESSAGE = `Your task is to predict the line number where the developer is most likely to make their next edit. If you jump in the current file, just output the line number. If you want to jump to another file, output the filepath (relative to workspace root), colon, then line number. If you don't think anywhere is a good next line jump target, just output the current line number of the cursor. Make sure to output no explanation, reasoning, extra spaces, etc.`;
+
 export class XtabNextCursorPredictor {
 
 	private isDisabled: boolean;
@@ -83,11 +92,15 @@ export class XtabNextCursorPredictor {
 	}
 
 
-	public async predictNextCursorPosition(promptPieces: PromptPieces, parentTracer: ILogger, telemetryBuilder: StatelessNextEditTelemetryBuilder | undefined, cancellationToken: CancellationToken): Promise<Result<CursorJumpPrediction, Error>> {
-
-		const tracer = parentTracer.createSubLogger('predictNextCursorPosition');
-
-		const systemMessage = `Your task is to predict the line number where the developer is most likely to make their next edit. If you jump in the current file, just output the line number. If you want to jump to another file, output the filepath (relative to workspace root), colon, then line number. If you don't think anywhere is a good next line jump target, just output the current line number of the cursor. Make sure to output no explanation, reasoning, extra spaces, etc.`;
+	/**
+	 * Build the chat messages and `keptRange` for the cursor-prediction prompt
+	 * without making any network call. Extracted from {@link predictNextCursorPosition}
+	 * so that test/datagen tooling can capture the exact production prompt
+	 * (and the kept-line range used to validate the assistant's response)
+	 * without spinning up a mock endpoint.
+	 */
+	public buildCursorPredictionPrompt(promptPieces: PromptPieces): Result<{ messages: Raw.ChatMessage[]; keptRange: OffsetRange }, Error> {
+		const systemMessage = NEXT_CURSOR_PREDICTION_SYSTEM_MESSAGE;
 
 		const maxTokens = this.configService.getExperimentBasedConfig(ConfigKey.Advanced.InlineEditsNextCursorPredictionCurrentFileMaxTokens, this.expService);
 
@@ -113,7 +126,6 @@ export class XtabNextCursorPredictor {
 		);
 
 		if (currentFileContentR.isError()) {
-			tracer.trace(`Failed to construct tagged file: ${currentFileContentR.err}`);
 			return Result.fromString(currentFileContentR.err);
 		}
 
@@ -163,10 +175,26 @@ export class XtabNextCursorPredictor {
 			userMsg: userMessage
 		});
 
-		telemetryBuilder?.setCursorJumpPrompt(messages);
+		return Result.ok({ messages, keptRange: clippedTaggedCurrentDoc.keptRange });
+	}
+
+
+	public async predictNextCursorPosition(promptPieces: PromptPieces, tracing: RequestTracingContext, cancellationToken: CancellationToken): Promise<Result<CursorJumpPrediction, Error>> {
+
+		const tracer = tracing.tracer.createSubLogger('predictNextCursorPosition');
+
+		const promptR = this.buildCursorPredictionPrompt(promptPieces);
+		if (promptR.isError()) {
+			tracer.trace(`Failed to construct tagged file: ${promptR.err.message}`);
+			return Result.fromString(promptR.err.message);
+		}
+		const { messages, keptRange } = promptR.val;
+
+		tracing.telemetry.setCursorJumpPrompt(messages);
+		tracing.logContext.setCursorJumpPrompt(messages, keptRange);
 
 		const modelName = this.determineModelName();
-		telemetryBuilder?.setCursorJumpModelName(modelName);
+		tracing.telemetry.setCursorJumpModelName(modelName);
 
 		const resolvedEndpoint = await this.resolveEndpoint(modelName, tracer);
 		if (!resolvedEndpoint) {
@@ -208,9 +236,9 @@ export class XtabNextCursorPredictor {
 		}
 
 		try {
-			telemetryBuilder?.setCursorJumpResponse(response.value);
+			tracing.telemetry.setCursorJumpResponse(response.value);
 			const trimmed = response.value.trim();
-			return this.parseResponse(trimmed, clippedTaggedCurrentDoc.keptRange);
+			return this.parseResponse(trimmed, keptRange);
 		} catch (err: unknown) {
 			tracer.trace(`Failed to parse predicted line number from response '${response.value}': ${err}`);
 			return Result.fromString(`failedToParseLine:"${response.value}". Error ${ErrorUtils.fromUnknown(err).message}`);

extensions/copilot/src/extension/xtab/node/xtabProvider.ts

@@ -98,7 +98,7 @@ export interface ModelConfig extends xtabPromptOptions.PromptOptions {
 	modelName: string | undefined;
 }
 
-interface RequestTracingContext {
+export interface RequestTracingContext {
 	tracer: ILogger;
 	logContext: InlineEditRequestLogContext;
 	telemetry: StatelessNextEditTelemetryBuilder;
@@ -1155,7 +1155,7 @@ export class XtabProvider implements IStatelessNextEditProvider {
 			return new NoNextEditReason.GotCancelled('beforeNextCursorPredictionFetchUserTyped');
 		}
 
-		const nextCursorLineR = await this.nextCursorPredictor.predictNextCursorPosition(promptPieces, tracer, telemetry, cancellationToken);
+		const nextCursorLineR = await this.nextCursorPredictor.predictNextCursorPosition(promptPieces, tracing, cancellationToken);
 
 		if (cancellationToken.isCancellationRequested) {
 			return new NoNextEditReason.GotCancelled('afterNextCursorPredictionFetch');

extensions/copilot/src/extension/xtab/test/node/xtabNextCursorPredictor.spec.ts

@@ -15,6 +15,8 @@
 import { AggressivenessLevel, DEFAULT_OPTIONS, PromptOptions } from '../../../../platform/inlineEdits/common/dataTypes/xtabPromptOptions';
 import { StatelessNextEditDocument } from '../../../../platform/inlineEdits/common/statelessNextEditProvider';
 import { TestLanguageDiagnosticsService } from '../../../../platform/languages/common/testLanguageDiagnosticsService';
+import { InlineEditRequestLogContext } from '../../../../platform/inlineEdits/common/inlineEditLogContext';
+import { StatelessNextEditTelemetryBuilder } from '../../../../platform/inlineEdits/common/statelessNextEditProvider';
 import { ILogger } from '../../../../platform/log/common/logService';
 import { ITestingServicesAccessor } from '../../../../platform/test/node/services';
 import { TestLogService } from '../../../../platform/testing/common/testLogService';
@@ -31,11 +33,20 @@
 import { PromptPieces } from '../../common/promptCrafting';
 import { CurrentDocument } from '../../common/xtabCurrentDocument';
 import { XtabNextCursorPredictor } from '../../node/xtabNextCursorPredictor';
+import type { RequestTracingContext } from '../../node/xtabProvider';
 
 function createTestLogger(): ILogger {
 	return new TestLogService();
 }
 
+function createTestTracingContext(tracer: ILogger): RequestTracingContext {
+	return {
+		tracer,
+		logContext: new InlineEditRequestLogContext('test', 0, undefined),
+		telemetry: new StatelessNextEditTelemetryBuilder('test-request'),
+	};
+}
+
 function computeTokens(s: string): number {
 	return Math.ceil(s.length / 4);
 }
@@ -130,7 +141,7 @@
 			});
 
 			// Make a prediction request - should fail with NotFound
-			const result = await predictor.predictNextCursorPosition(promptPieces, tracer, undefined, CancellationToken.None);
+			const result = await predictor.predictNextCursorPosition(promptPieces, createTestTracingContext(tracer), CancellationToken.None);
 
 			expect(result.isError()).toBe(true);
 			if (result.isError()) {
@@ -155,7 +166,7 @@
 			});
 
 			// First call - triggers disabling
-			await predictor.predictNextCursorPosition(promptPieces, tracer, undefined, CancellationToken.None);
+			await predictor.predictNextCursorPosition(promptPieces, createTestTracingContext(tracer), CancellationToken.None);
 
 			// Verify disabled
 			expect(predictor.determineEnablement()).toBeUndefined();
@@ -191,7 +202,7 @@
 			});
 
 			// Make a prediction request - should fail but not disable
-			const result = await predictor.predictNextCursorPosition(promptPieces, tracer, undefined, CancellationToken.None);
+			const result = await predictor.predictNextCursorPosition(promptPieces, createTestTracingContext(tracer), CancellationToken.None);
 
 			expect(result.isError()).toBe(true);
 			if (result.isError()) {
@@ -217,7 +228,7 @@
 				resolvedModel: 'test-model'
 			});
 
-			const result = await predictor.predictNextCursorPosition(promptPieces, tracer, undefined, CancellationToken.None);
+			const result = await predictor.predictNextCursorPosition(promptPieces, createTestTracingContext(tracer), CancellationToken.None);
 
 			expect(result.isOk()).toBe(true);
 			if (result.isOk()) {
@@ -242,7 +253,7 @@
 				resolvedModel: 'test-model'
 			});
 
-			const result = await predictor.predictNextCursorPosition(promptPieces, tracer, undefined, CancellationToken.None);
+			const result = await predictor.predictNextCursorPosition(promptPieces, createTestTracingContext(tracer), CancellationToken.None);
 
 			expect(result.isOk()).toBe(true);
 			if (result.isOk()) {

extensions/copilot/src/platform/inlineEdits/common/inlineEditLogContext.ts

@@ -11,6 +11,7 @@ import { isCancellationError } from '../../../util/vs/base/common/errors';
 import { Emitter, Event } from '../../../util/vs/base/common/event';
 import { ThemeIcon } from '../../../util/vs/base/common/themables';
 import { SerializedLineEdit } from '../../../util/vs/editor/common/core/edits/lineEdit';
+import { OffsetRange } from '../../../util/vs/editor/common/core/ranges/offsetRange';
 import { SerializedEdit } from './dataTypes/editUtils';
 import { FetchCancellationError } from './dataTypes/fetchCancellationError';
 import { LanguageContextResponse, SerializedContextResponse, serializeLanguageContext } from './dataTypes/languageContext';
@@ -384,6 +385,31 @@ export class InlineEditRequestLogContext {
 		this.fireDidChange();
 	}
 
+	/**
+	 * Raw chat messages used to construct the cursor-jump (next-cursor-line
+	 * prediction) prompt, and the document-line offset range the model can
+	 * reference in its response. Stored here so in-process debug / datagen
+	 * tooling can read them back via the same log-context vehicle as the
+	 * xtab prompt (`rawMessages`). Never emitted to telemetry sinks — the
+	 * `rawMessages` can contain full prompt content (source code).
+	 */
+	private _cursorJumpRawMessages: Raw.ChatMessage[] | undefined = undefined;
+	private _cursorJumpKeptRange: OffsetRange | undefined = undefined;
+
+	get cursorJumpRawMessages(): Raw.ChatMessage[] | undefined {
+		return this._cursorJumpRawMessages;
+	}
+
+	get cursorJumpKeptRange(): OffsetRange | undefined {
+		return this._cursorJumpKeptRange;
+	}
+
+	setCursorJumpPrompt(messages: Raw.ChatMessage[], keptRange: OffsetRange) {
+		this._cursorJumpRawMessages = messages;
+		this._cursorJumpKeptRange = keptRange;
+		this.fireDidChange();
+	}
+
 	private _outcome: LogContextOutcome = 'pending';
 
 	/**

extensions/copilot/test/base/simulationOptions.ts

@@ -9,11 +9,23 @@ import { CacheMode } from './simulationContext';
 /** Number of runs that are stored in baseline.json */
 export const BASELINE_RUN_COUNT = 10;
 
+export enum NesDatagenSampleTask {
+	Xtab = 'xtab',
+	CursorSameFile = 'cursor-same-file',
+	CursorCrossFile = 'cursor-cross-file',
+	CursorBoth = 'cursor-both',
+}
+
 export type NesDatagen = {
 	readonly input: string;
 	readonly output: string | undefined;
 	readonly rowOffset: number;
 	readonly workerMode: boolean;
+	readonly sampleTask: NesDatagenSampleTask;
+	/** Minimum same-file lines above the request cursor for a move to count as a jump. */
+	readonly sameFileJumpMinAbove: number;
+	/** Minimum same-file lines below the request cursor for a move to count as a jump. */
+	readonly sameFileJumpMinBelow: number;
 };
 
 export class SimulationOptions {
@@ -175,6 +187,9 @@ export class SimulationOptions {
 				output: argv['out'],
 				rowOffset: typeof argv['row-offset'] === 'number' ? argv['row-offset'] : 0,
 				workerMode: boolean(argv['worker'], false),
+				sampleTask: SimulationOptions.validateSampleTask(argv['sample-task']),
+				sameFileJumpMinAbove: typeof argv['same-file-jump-min-above'] === 'number' ? argv['same-file-jump-min-above'] : 2,
+				sameFileJumpMinBelow: typeof argv['same-file-jump-min-below'] === 'number' ? argv['same-file-jump-min-below'] : 5,
 			}
 			: undefined;
 
@@ -252,6 +267,14 @@ export class SimulationOptions {
 			`  --input                            Path to a JSON or JSON Lines file with training data recordings (required)`,
 			`                                     Format is inferred from the extension: .jsonl/.ndjson → JSON Lines, otherwise JSON array`,
 			`  --out                              Output path for the JSON Lines file. Default: <input-path>_output.jsonl`,
+			`  --sample-task                      Which target to generate (default: xtab)`,
+			`                                     Values: xtab, cursor-same-file, cursor-cross-file, cursor-both`,
+			`                                       xtab               → edit-prediction sample (assistant = an edit)`,
+			`                                       cursor-same-file   → next-cursor-line sample restricted to the active file`,
+			`                                       cursor-cross-file  → next-cursor-line sample for a jump to another file`,
+			`                                       cursor-both        → tries same-file first, falls back to cross-file (one sample per row)`,
+			`  --same-file-jump-min-above         Minimum lines above request cursor for a same-file move to count as a jump (default: 2)`,
+			`  --same-file-jump-min-below         Minimum lines below request cursor for a same-file move to count as a jump (default: 5)`,
 			``,
 			`Global options (placed before 'nes-datagen'):`,
 			`  --config-file                      Path to a JSON config file (required for nes-datagen)`,
@@ -264,6 +287,9 @@ export class SimulationOptions {
 			`Examples:`,
 			`  npm run simulate -- --config-file=config.json nes-datagen --input=data.json`,
 			`  npm run simulate -- --config-file=config.json --parallelism=10 --verbose nes-datagen --input=data.json`,
+			`  npm run simulate -- --config-file=config.json nes-datagen --input=data.json --sample-task=cursor-same-file`,
+			`  npm run simulate -- --config-file=config.json nes-datagen --input=data.json --sample-task=cursor-cross-file`,
+			`  npm run simulate -- --config-file=config.json nes-datagen --input=data.json --sample-task=cursor-both --same-file-jump-min-above=8 --same-file-jump-min-below=8`,
 			``,
 		].join('\n'));
 	}
@@ -301,6 +327,20 @@ export class SimulationOptions {
 			throw new Error(`--nesUrl must be provided when --nesApiKey is set`);
 		}
 	}
+
+	private static validateSampleTask(value: unknown): NesDatagenSampleTask {
+		if (value === undefined || value === null) {
+			return NesDatagenSampleTask.Xtab;
+		}
+		if (typeof value !== 'string') {
+			throw new Error(`--sample-task must be a string, but got: ${typeof value}`);
+		}
+		const allowed = Object.values(NesDatagenSampleTask) as string[];
+		if (!allowed.includes(value)) {
+			throw new Error(`--sample-task must be one of [${allowed.join(', ')}], but got: ${value}`);
+		}
+		return value as NesDatagenSampleTask;
+	}
 }
 
 function cliOptionsToWellKnownEmbeddingsType(model: string | undefined): EmbeddingType | undefined {