removing stale compaction item

Author: dileepyavan

src/platform/endpoint/node/responsesApi.ts

@@ -30,14 +30,15 @@ import { getStatefulMarkerAndIndex } from '../common/statefulMarkerContainer';
 import { rawPartAsThinkingData } from '../common/thinkingDataContainer';
 
 export function getResponsesApiCompactionThreshold(configService: IConfigurationService, expService: IExperimentationService, endpoint: IChatEndpoint): number | undefined {
-	const contextManagementEnabled = configService.getExperimentBasedConfig(ConfigKey.ResponsesApiContextManagementEnabled, expService) && !modelsWithoutResponsesContextManagement.has(endpoint.family);
+	const contextManagementEnabled = !modelsWithoutResponsesContextManagement.has(endpoint.family);
 	if (!contextManagementEnabled) {
 		return undefined;
 	}
 
-	return endpoint.modelMaxPromptTokens > 0
-		? Math.floor(endpoint.modelMaxPromptTokens * 0.9)
-		: 50000;
+	// return endpoint.modelMaxPromptTokens > 0
+	// 	? Math.floor(endpoint.modelMaxPromptTokens * 0.9)
+	// 	: 50000;
+	return 10000;
 }
 
 export function createResponsesRequestBody(accessor: ServicesAccessor, options: ICreateEndpointBodyOptions, model: string, endpoint: IChatEndpoint): IEndpointBody {
@@ -49,7 +50,7 @@ export function createResponsesRequestBody(accessor: ServicesAccessor, options:
 
 	const body: IEndpointBody = {
 		model,
-		...rawMessagesToResponseAPI(model, options.messages, !!options.ignoreStatefulMarker, compactThreshold !== undefined),
+		...rawMessagesToResponseAPI(model, options.messages, !!options.ignoreStatefulMarker),
 		stream: true,
 		tools: options.requestOptions?.tools?.map((tool): OpenAI.Responses.FunctionTool & OpenAiResponsesFunctionTool => ({
 			...tool.function,
@@ -124,19 +125,32 @@ interface ResponseOutputItemWithPhase {
 	phase?: string;
 }
 
-function rawMessagesToResponseAPI(modelId: string, messages: readonly Raw.ChatMessage[], ignoreStatefulMarker: boolean, compactionEnabled: boolean): { input: OpenAI.Responses.ResponseInputItem[]; previous_response_id?: string } {
+interface LatestCompactionOutput {
+	readonly item: OpenAIContextManagementResponse;
+	readonly outputIndex: number;
+}
+
+function rawMessagesToResponseAPI(modelId: string, messages: readonly Raw.ChatMessage[], ignoreStatefulMarker: boolean): { input: OpenAI.Responses.ResponseInputItem[]; previous_response_id?: string } {
 	const latestCompactionMessageIndex = getLatestCompactionMessageIndex(messages);
+	const latestCompactionMessage = latestCompactionMessageIndex !== undefined ? createCompactionRoundTripMessage(messages[latestCompactionMessageIndex]) : undefined;
 	const statefulMarkerAndIndex = !ignoreStatefulMarker && getStatefulMarkerAndIndex(modelId, messages);
 
 	let previousResponseId: string | undefined;
 	if (statefulMarkerAndIndex) {
 		previousResponseId = statefulMarkerAndIndex.statefulMarker;
-		if (!compactionEnabled) {
-			messages = messages.slice(statefulMarkerAndIndex.index + 1);
-		}
-	}
 
-	if (latestCompactionMessageIndex !== undefined) {
+		// Requests that resume from previous_response_id send only post-marker history,
+		// but they still need the latest compaction item even when that item predates
+		// the marker. This keeps both websocket and non-websocket traffic aligned.
+		messages = messages.slice(statefulMarkerAndIndex.index + 1);
+		if (latestCompactionMessageIndex !== undefined) {
+			if (latestCompactionMessageIndex > statefulMarkerAndIndex.index) {
+				messages = messages.slice(latestCompactionMessageIndex - (statefulMarkerAndIndex.index + 1));
+			} else if (latestCompactionMessage) {
+				messages = [latestCompactionMessage, ...messages];
+			}
+		}
+	} else if (latestCompactionMessageIndex !== undefined) {
 		messages = messages.slice(latestCompactionMessageIndex);
 	}
 
@@ -200,6 +214,22 @@ function rawMessagesToResponseAPI(modelId: string, messages: readonly Raw.ChatMe
 	return { input, previous_response_id: previousResponseId };
 }
 
+function createCompactionRoundTripMessage(message: Raw.ChatMessage): Raw.ChatMessage | undefined {
+	if (message.role !== Raw.ChatRole.Assistant) {
+		return undefined;
+	}
+
+	const content = message.content.filter(part => part.type === Raw.ChatCompletionContentPartKind.Opaque && rawPartAsCompactionData(part));
+	if (!content.length) {
+		return undefined;
+	}
+
+	return {
+		role: Raw.ChatRole.Assistant,
+		content,
+	};
+}
+
 function getLatestCompactionMessageIndex(messages: readonly Raw.ChatMessage[]): number | undefined {
 	for (let idx = messages.length - 1; idx >= 0; idx--) {
 		const message = messages[idx];
@@ -452,6 +482,39 @@ function responseFunctionOutputToRawContents(output: string | OpenAI.Responses.R
 	return coalesce(output.map(responseContentToRawContent));
 }
 
+function isCompactionOutputItem(item: OpenAI.Responses.ResponseOutputItem): boolean {
+	return item.type.toString() === openAIContextManagementCompactionType;
+}
+
+function getLatestCompactionOutput(output: OpenAI.Responses.ResponseOutputItem[], preferredOutputIndex: number | undefined): LatestCompactionOutput | undefined {
+	let latestCompactionOutput: LatestCompactionOutput | undefined;
+	for (let idx = output.length - 1; idx >= 0; idx--) {
+		const item = output[idx];
+		if (isCompactionOutputItem(item)) {
+			latestCompactionOutput = { item: item as unknown as OpenAIContextManagementResponse, outputIndex: idx };
+			break;
+		}
+	}
+
+	if (preferredOutputIndex !== undefined) {
+		const preferredItem = output[preferredOutputIndex];
+		if (preferredItem && isCompactionOutputItem(preferredItem) && (!latestCompactionOutput || preferredOutputIndex >= latestCompactionOutput.outputIndex)) {
+			return { item: preferredItem as unknown as OpenAIContextManagementResponse, outputIndex: preferredOutputIndex };
+		}
+	}
+
+	return latestCompactionOutput;
+}
+
+function keepLatestCompactionOutput(output: OpenAI.Responses.ResponseOutputItem[], preferredOutputIndex: number | undefined): OpenAI.Responses.ResponseOutputItem[] {
+	const latestCompactionOutput = getLatestCompactionOutput(output, preferredOutputIndex);
+	if (!latestCompactionOutput) {
+		return output;
+	}
+
+	return output.filter((item, idx) => !isCompactionOutputItem(item) || idx === latestCompactionOutput.outputIndex);
+}
+
 export async function processResponseFromChatEndpoint(instantiationService: IInstantiationService, telemetryService: ITelemetryService, logService: ILogService, response: Response, expectedNumChoices: number, finishCallback: FinishedCallback, telemetryData: TelemetryData, compactionThreshold?: number): Promise<AsyncIterableObject<ChatCompletion>> {
 	return new AsyncIterableObject<ChatCompletion>(async feed => {
 		const requestId = response.headers.get('X-Request-ID') ?? generateUuid();
@@ -499,6 +562,8 @@ export class OpenAIResponsesProcessor {
 	private textAccumulator: string = '';
 	private hasReceivedReasoningSummary = false;
 	private sawCompactionMessage = false;
+	private latestCompactionOutputIndex: number | undefined;
+	private latestCompactionItem: OpenAIContextManagementResponse | undefined;
 	/** Maps output_index to { name, callId, arguments } for streaming tool call updates */
 	private readonly toolCallInfo = new Map<number, { name: string; callId: string; arguments: string }>();
 
@@ -560,6 +625,12 @@ export class OpenAIResponsesProcessor {
 			case 'response.output_item.done':
 				if (chunk.item.type.toString() === openAIContextManagementCompactionType) {
 					const compactionItem = chunk.item as unknown as OpenAIContextManagementResponse;
+					if (this.latestCompactionOutputIndex !== undefined && chunk.output_index < this.latestCompactionOutputIndex) {
+						return;
+					}
+
+					this.latestCompactionOutputIndex = chunk.output_index;
+					this.latestCompactionItem = compactionItem;
 					this.sawCompactionMessage = true;
 					return onProgress({
 						text: '',
@@ -617,33 +688,58 @@ export class OpenAIResponsesProcessor {
 						id: chunk.item_id
 					}
 				});
-			case 'response.completed':
+			case 'response.completed': {
+				const normalizedOutput = keepLatestCompactionOutput(chunk.response.output, this.latestCompactionOutputIndex);
+				const latestCompactionOutput = getLatestCompactionOutput(normalizedOutput, this.latestCompactionOutputIndex);
+				const latestCompactionItem = latestCompactionOutput?.item;
+				const previousCompactionItem = this.latestCompactionItem;
+				if (latestCompactionItem) {
+					this.sawCompactionMessage = true;
+					this.latestCompactionOutputIndex = latestCompactionOutput.outputIndex;
+				}
+
+				const shouldEmitResolvedCompaction = latestCompactionItem && (
+					!previousCompactionItem ||
+					previousCompactionItem.id !== latestCompactionItem.id ||
+					previousCompactionItem.encrypted_content !== latestCompactionItem.encrypted_content
+				);
+				if (latestCompactionItem) {
+					this.latestCompactionItem = latestCompactionItem;
+				}
 				if (this.compactionThreshold !== undefined && this.sawCompactionMessage) {
+					const promptTokens = chunk.response.usage?.input_tokens ?? 0;
+					const totalTokens = chunk.response.usage?.total_tokens ?? 0;
 					sendResponsesApiCompactionTelemetry(this.telemetryService, {
 						outcome: 'compaction_returned',
 						headerRequestId: this.requestId,
 						gitHubRequestId: this.ghRequestId,
 						model: chunk.response.model,
 					}, {
 						compactThreshold: this.compactionThreshold,
-						promptTokens: chunk.response.usage?.input_tokens ?? 0,
-						totalTokens: chunk.response.usage?.total_tokens ?? 0,
+						promptTokens,
+						totalTokens,
 					});
 					this.logService.debug(`[responsesAPI_compaction] Compaction enabled. headerRequestId=${this.requestId}`);
 				} else if (this.compactionThreshold !== undefined && (chunk.response.usage?.input_tokens ?? 0) >= this.compactionThreshold) {
+					const promptTokens = chunk.response.usage?.input_tokens ?? 0;
+					const totalTokens = chunk.response.usage?.total_tokens ?? 0;
 					sendResponsesApiCompactionTelemetry(this.telemetryService, {
 						outcome: 'threshold_met_no_compaction',
 						headerRequestId: this.requestId,
 						gitHubRequestId: this.ghRequestId,
 						model: chunk.response.model,
 					}, {
 						compactThreshold: this.compactionThreshold,
-						promptTokens: chunk.response.usage?.input_tokens ?? 0,
-						totalTokens: chunk.response.usage?.total_tokens ?? 0,
+						promptTokens,
+						totalTokens,
 					});
-					this.logService.debug(`[responsesAPI_compaction] Compaction enabled but context not compacted after threshold was met. headerRequestId=${this.requestId}`);
+					this.logService.debug(`[responsesAPI_compaction] Compaction enabled but context not compacted after threshold was met. headerRequestId=${this.requestId}, gitHubRequestId=${this.ghRequestId}, promptTokens=${promptTokens}, totalTokens=${totalTokens}`);
 				}
-				onProgress({ text: '', statefulMarker: chunk.response.id });
+				onProgress({
+					text: '',
+					statefulMarker: chunk.response.id,
+					contextManagement: shouldEmitResolvedCompaction ? latestCompactionItem : undefined,
+				});
 				return {
 					blockFinished: true,
 					choiceIndex: 0,
@@ -667,7 +763,7 @@ export class OpenAIResponsesProcessor {
 					finishReason: FinishedCompletionReason.Stop,
 					message: {
 						role: Raw.ChatRole.Assistant,
-						content: chunk.response.output.map((item): Raw.ChatCompletionContentPart | undefined => {
+						content: normalizedOutput.map((item): Raw.ChatCompletionContentPart | undefined => {
 							if (item.type === 'message') {
 								return { type: Raw.ChatCompletionContentPartKind.Text, text: item.content.map(c => c.type === 'output_text' ? c.text : c.refusal).join('') };
 							} else if (item.type === 'image_generation_call' && item.result) {
@@ -676,6 +772,7 @@ export class OpenAIResponsesProcessor {
 						}).filter(isDefined),
 					}
 				};
+			}
 		}
 	}
 }