[content] src/extension/completions-core/vscode-node/lib/src/postInsertion.ts (typescript)
import { CopilotNamedAnnotationList } from './openai/stream';
import { contextIndentationFromText, indentationBlockFinished } from './prompt/parseBlock';
import { Prompt, extractPrompt } from './prompt/prompt';
import { fetchCitations } from './snippy/handlePostInsertion';
import { editDistance, lexEditDistance } from './suggestions/editDistance';
import { SuggestionStatus, computeCompletionText } from './suggestions/partialSuggestions';
import { TelemetryStore, TelemetryWithExp, telemetry, telemetryCatch } from './telemetry';
import { TextDocumentManager } from './textDocumentManager';
import { ICompletionsPromiseQueueService } from './util/promiseQueue';
import { ICompletionsRuntimeModeService } from './util/runtimeMode';

const postInsertionLogger = new Logger('postInsertion');

type Timeout = {
	seconds: number;
	captureCode: boolean;
	captureRejection: boolean;
};
// windows for telemetry checks, in seconds
// captureCode = capture the code after acceptance,
// captureRejection = capture the code after rejection
const captureTimeouts: Timeout[] = [
	{ seconds: 15, captureCode: false, captureRejection: false },
	{ seconds: 30, captureCode: true, captureRejection: true },
	{ seconds: 120, captureCode: false, captureRejection: false },
	{ seconds: 300, captureCode: false, captureRejection: false },
	{ seconds: 600, captureCode: false, captureRejection: false },
];

// No. of chars before/after insertion point to look for the completion
const stillInCodeNearMargin = 50;
const stillInCodeFarMargin = 1500;

// If lex edit distance is below this fraction of completion length it is considered
// in the code
const stillInCodeFraction = 0.5;

// Number of characters captured after the insertion point.
// Used only if we couldn't detect termination point with indent-based parsing.
const captureCodeMargin = 500;

const postInsertConfiguration: {
	triggerPostInsertionSynchroneously: boolean;
	captureCode: boolean;
	captureRejection: boolean;
} = {
	triggerPostInsertionSynchroneously: false,
	captureCode: false,
	captureRejection: false,
};

async function captureCode(
	accessor: ServicesAccessor,
	uri: string,
	completionTelemetry: TelemetryWithExp,
	offset: number,
	suffixOffset?: number
): Promise<{ prompt: Prompt; capturedCode: string; terminationOffset: number }> {
	const ctx = accessor.get(ICompletionsContextService);
	const instantiationService = accessor.get(IInstantiationService);
	const logTarget = ctx.get(LogTarget);
	const result = await ctx.get(FileReader).getOrReadTextDocumentWithFakeClientProperties({ uri });
	if (result.status !== 'valid') {
		postInsertionLogger.info(logTarget, `Could not get document for ${uri}. Maybe it was closed by the editor.`);
		return {
			prompt: {
				prefix: '',
				suffix: '',
				isFimEnabled: false,
			},
			capturedCode: '',
			terminationOffset: 0,
		};
	}
	const document = result.document;
	const documentText = document.getText();
	const documentTextBefore = documentText.substring(0, offset);
	const position = document.positionAt(offset);

	// Treat the code before offset as the hypothetical prompt
	const hypotheticalPromptResponse = await instantiationService.invokeFunction(extractPrompt,
		completionTelemetry.properties.headerRequestId,
		createCompletionState(document, position),
		completionTelemetry
	);
	const hypotheticalPrompt =
		hypotheticalPromptResponse.type === 'prompt'
			? hypotheticalPromptResponse.prompt
			: {
				prefix: documentTextBefore,
				suffix: '',
				isFimEnabled: false,
			}; // TODO(eaftan): Pass an actual suffix when we're ready to support it

	if (hypotheticalPrompt.isFimEnabled && suffixOffset !== undefined) {
		// With FIM enabled, we can exactly determine capturedCode, suffix and prefix by propertly initialized trackers. No need to guess.
		const capturedCode = documentText.substring(offset, suffixOffset);
		hypotheticalPrompt.suffix = documentText.substring(suffixOffset);

		return { prompt: hypotheticalPrompt, capturedCode, terminationOffset: 0 };
	} else {
		//Everything after the insertion point is hypothetical response we could get from AI
		const hypotheticalResponse = documentText.substring(offset);

		//Try to find the termination offset in the hypothetical response using indentation based parsing
		const contextIndent = contextIndentationFromText(documentTextBefore, offset, document.detectedLanguageId);
		const indentTerminationFunction = indentationBlockFinished(contextIndent, undefined);
		const terminationResult = indentTerminationFunction(hypotheticalResponse);

		//If we could detect termination of the indentation block, capture 2x length of detected suggestion
		//Otherwise capture a lot of characters
		const maxOffset = Math.min(
			documentText.length,
			offset + (terminationResult ? terminationResult * 2 : captureCodeMargin)
		);

		const capturedCode = documentText.substring(offset, maxOffset);

		return { prompt: hypotheticalPrompt, capturedCode, terminationOffset: terminationResult ?? -1 };
	}
}

export function postRejectionTasks(
	accessor: ServicesAccessor,
	insertionCategory: PostInsertionCategory,
	insertionOffset: number,
	uri: string,
	completions: { completionText: string; completionTelemetryData: TelemetryWithExp }[]
) {
	const logTarget = accessor.get(ICompletionsContextService).get(LogTarget);
	const instantiationService = accessor.get(IInstantiationService);
	const telemetryService = accessor.get(ICompletionsTelemetryService);
	const promiseQueueService = accessor.get(ICompletionsPromiseQueueService);

	//Send `.rejected` telemetry event for each rejected completion
	completions.forEach(({ completionText, completionTelemetryData }) => {
		postInsertionLogger.debug(
			logTarget,
			`${insertionCategory}.rejected choiceIndex: ${completionTelemetryData.properties.choiceIndex}`
		);
		instantiationService.invokeFunction(telemetryRejected, insertionCategory, completionTelemetryData);
	});
	const positionTracker = instantiationService.createInstance(ChangeTracker, uri, insertionOffset - 1);
	const suffixTracker = instantiationService.createInstance(ChangeTracker, uri, insertionOffset);

	const checkInCode = async (t: Timeout) => {
		postInsertionLogger.debug(
			logTarget,
			`Original offset: ${insertionOffset}, Tracked offset: ${positionTracker.offset}`
		);
		const { completionTelemetryData } = completions[0];

		const { prompt, capturedCode, terminationOffset } = await instantiationService.invokeFunction(captureCode,
			uri,
			completionTelemetryData,
			positionTracker.offset + 1,
			suffixTracker.offset
		);

		const promptTelemetry = {
			hypotheticalPromptJson: JSON.stringify({ prefix: prompt.prefix, context: prompt.context }),
			hypotheticalPromptSuffixJson: JSON.stringify(prompt.suffix),
		};

		const customTelemetryData = completionTelemetryData.extendedBy(
			{
				...promptTelemetry,
				capturedCodeJson: JSON.stringify(capturedCode),
			},
			{
				timeout: t.seconds,
				insertionOffset: insertionOffset,
				trackedOffset: positionTracker.offset,
				terminationOffsetInCapturedCode: terminationOffset,
			}
		);
		postInsertionLogger.debug(
			logTarget,
			`${insertionCategory}.capturedAfterRejected choiceIndex: ${completionTelemetryData.properties.choiceIndex}`,
			customTelemetryData
		);
		instantiationService.invokeFunction(telemetry, insertionCategory + '.capturedAfterRejected', customTelemetryData, TelemetryStore.Enhanced);
	};
	// Capture the code typed after we detected that completion was rejected,
	// Uses first displayed completion as the source/seed of telemetry information.
	captureTimeouts
		.filter(t => t.captureRejection)
		.map(t =>
			positionTracker.push(
				telemetryCatch(telemetryService, promiseQueueService, () => checkInCode(t), 'postRejectionTasks'),
				t.seconds * 1000
			)
		);
}

export function postInsertionTasks(
	accessor: ServicesAccessor,
	insertionCategory: PostInsertionCategory,
	completionText: string,
	insertionOffset: number,
	uri: string,
	telemetryData: TelemetryWithExp,
	suggestionStatus: SuggestionStatus,
	copilotAnnotations?: CopilotNamedAnnotationList
) {
	const logTarget = accessor.get(ICompletionsContextService).get(LogTarget);
	const instantiationService = accessor.get(IInstantiationService);
	const promiseQueueService = accessor.get(ICompletionsPromiseQueueService);
	const telemetryService = accessor.get(ICompletionsTelemetryService);
	const runtimeModeService = accessor.get(ICompletionsRuntimeModeService);

	const telemetryDataWithStatus = telemetryData.extendedBy(
		{
			compType: suggestionStatus.compType,
		},
		{
			compCharLen: suggestionStatus.acceptedLength,
			numLines: suggestionStatus.acceptedLines,
		}
	);
	// send ".accepted" telemetry
	postInsertionLogger.debug(
		logTarget,
		`${insertionCategory}.accepted choiceIndex: ${telemetryDataWithStatus.properties.choiceIndex}`
	);
	instantiationService.invokeFunction(telemetryAccepted, insertionCategory, telemetryDataWithStatus);

	const fullCompletionText = completionText;
	completionText = computeCompletionText(completionText, suggestionStatus);
	const trimmedCompletion = completionText.trim();
	const tracker = instantiationService.createInstance(ChangeTracker, uri, insertionOffset);
	const suffixTracker = instantiationService.createInstance(ChangeTracker, uri, insertionOffset + completionText.length);