Fix cost calc + context size compaction (#321980)

* Fix cost calc + context size compaction

* Address comments

extensions/copilot/src/extension/intents/node/agentIntent.ts

@@ -82,6 +82,34 @@ function isResponsesCompactionContextManagementEnabled(endpoint: IChatEndpoint,
 		&& !modelsWithoutResponsesContextManagement.has(endpoint.family);
 }
 
+/**
+ * Applies the user's "Context Size" model-picker selection to the endpoint used
+ * for the agent's model requests.
+ *
+ * The picker offers two tiers — the model's default context max and its full
+ * native window (see `getContextSizeOptions`). For server-managed context (the
+ * Responses-API compaction path) the request endpoint's `modelMaxPromptTokens`
+ * is what drives the `compact_threshold` sent to the server. If the default
+ * tier is not propagated to the request endpoint, the server compacts against
+ * the model's full window and the stateful conversation grows far past the
+ * user's selection — billing them for the larger context. Mirrors the override
+ * applied on the `vscode.lm` path in `languageModelAccess.ts`.
+ *
+ * Only clamps when the selection is strictly smaller than the model window so
+ * the full tier ("Longer sessions without compaction") stays uncompacted.
+ *
+ * @internal - exported for testing
+ */
+export function applyContextSizeOverride(endpoint: IChatEndpoint, request: vscode.ChatRequest): IChatEndpoint {
+	const contextSize = request.modelConfiguration?.contextSize;
+	// Guard against non-positive / non-finite selections (e.g. 0, -1, NaN, Infinity):
+	// a non-positive token budget would produce an invalid endpoint configuration.
+	if (typeof contextSize === 'number' && Number.isFinite(contextSize) && contextSize > 0 && contextSize < endpoint.modelMaxPromptTokens) {
+		return endpoint.cloneWithTokenOverride(contextSize);
+	}
+	return endpoint;
+}
+
 /**
  * Returns true when the user explicitly referenced the todo tool (e.g. typed
  * `#todo` in their message) or a custom agent configuration includes it as a
@@ -613,7 +641,11 @@ export class AgentIntentInvocation extends EditCodeIntentInvocation implements I
 		@IAutomaticInstructionsCollector private readonly _automaticInstructionsCollector: IAutomaticInstructionsCollector,
 		@IAuthenticationService private readonly authenticationService: IAuthenticationService,
 	) {
-		super(intent, location, endpoint, request, intentOptions, instantiationService, codeMapperService, envService, promptPathRepresentationService, _endpointProvider, workspaceService, toolsService, configurationService, editLogService, commandService, telemetryService, notebookService, otelService);
+		// Apply the user's "Context Size" picker selection to the request endpoint
+		// so the server-managed compaction threshold (Responses API) is keyed to the
+		// selected tier rather than the model's full native window. See
+		// applyContextSizeOverride for the cost rationale.
+		super(intent, location, applyContextSizeOverride(endpoint, request), request, intentOptions, instantiationService, codeMapperService, envService, promptPathRepresentationService, _endpointProvider, workspaceService, toolsService, configurationService, editLogService, commandService, telemetryService, notebookService, otelService);
 	}
 
 	public override getAvailableTools(): Promise<vscode.LanguageModelToolInformation[]> {

extensions/copilot/src/extension/intents/node/test/contextSizeOverride.spec.ts

@@ -0,0 +1,57 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import type { ChatRequest } from 'vscode';
+import { describe, expect, test } from 'vitest';
+import { IChatEndpoint } from '../../../../platform/networking/common/networking';
+import { applyContextSizeOverride } from '../agentIntent';
+
+describe('applyContextSizeOverride', () => {
+	function createEndpoint(modelMaxPromptTokens: number): { endpoint: IChatEndpoint; clonedWith: number[] } {
+		const clonedWith: number[] = [];
+		const endpoint = {
+			modelMaxPromptTokens,
+			cloneWithTokenOverride(tokens: number): IChatEndpoint {
+				clonedWith.push(tokens);
+				return createEndpoint(tokens).endpoint;
+			},
+		} as unknown as IChatEndpoint;
+		return { endpoint, clonedWith };
+	}
+
+	function createRequest(contextSize?: unknown): ChatRequest {
+		return { modelConfiguration: contextSize === undefined ? undefined : { contextSize } } as unknown as ChatRequest;
+	}
+
+	test('clamps to the picked size when below the model window (default tier)', () => {
+		const { endpoint, clonedWith } = createEndpoint(400_000);
+		const result = applyContextSizeOverride(endpoint, createRequest(272_000));
+		expect(clonedWith).toEqual([272_000]);
+		expect(result.modelMaxPromptTokens).toBe(272_000);
+	});
+
+	test('leaves the endpoint untouched on the full tier (selection >= model window)', () => {
+		const { endpoint, clonedWith } = createEndpoint(400_000);
+		expect(applyContextSizeOverride(endpoint, createRequest(400_000))).toBe(endpoint);
+		expect(applyContextSizeOverride(endpoint, createRequest(500_000))).toBe(endpoint);
+		expect(clonedWith).toEqual([]);
+	});
+
+	test('does not clamp when context size is unset or non-numeric', () => {
+		const { endpoint, clonedWith } = createEndpoint(400_000);
+		expect(applyContextSizeOverride(endpoint, createRequest(undefined))).toBe(endpoint);
+		expect(applyContextSizeOverride(endpoint, createRequest('big'))).toBe(endpoint);
+		expect(clonedWith).toEqual([]);
+	});
+
+	test('does not clamp for non-positive or non-finite selections', () => {
+		const { endpoint, clonedWith } = createEndpoint(400_000);
+		expect(applyContextSizeOverride(endpoint, createRequest(0))).toBe(endpoint);
+		expect(applyContextSizeOverride(endpoint, createRequest(-1))).toBe(endpoint);
+		expect(applyContextSizeOverride(endpoint, createRequest(Number.NaN))).toBe(endpoint);
+		expect(applyContextSizeOverride(endpoint, createRequest(Number.POSITIVE_INFINITY))).toBe(endpoint);
+		expect(clonedWith).toEqual([]);
+	});
+});

extensions/copilot/src/extension/intents/node/toolCallingLoop.ts

@@ -180,6 +180,16 @@ export abstract class ToolCallingLoop<TOptions extends IToolCallingLoopOptions =
 	private lastHeaderRequestId: string | undefined;
 	private lastModelCallId: string | undefined;
 
+	/**
+	 * Running total of Copilot credits across every model call in the current
+	 * turn. Each fetch reports the credits for that single call, but the context
+	 * usage widget and `IChatModel.sessionCost` treat the per-request value as the
+	 * whole turn, so we accumulate here and emit the running total — mirroring the
+	 * cumulative credits the agent host reports. Reset on the first iteration of a
+	 * turn ({@link runOne} with `iterationNumber === 0`).
+	 */
+	private _accumulatedCopilotCredits: number | undefined;
+
 	/**
 	 * The full {@link ToolCallingLoopFetchOptions} from the most recent fetch.
 	 * Probes reuse this wholesale (overriding only `messages` and `finishedCb`)
@@ -1396,6 +1406,11 @@ export abstract class ToolCallingLoop<TOptions extends IToolCallingLoopOptions =
 
 	/** Runs a single iteration of the tool calling loop. */
 	public async runOne(outputStream: ChatResponseStream | undefined, iterationNumber: number, token: CancellationToken): Promise<IToolCallSingleResult> {
+		// The first iteration of a turn starts a fresh credit total. Resetting here
+		// (rather than only in run()) keeps runOne() correct when called standalone.
+		if (iterationNumber === 0) {
+			this._accumulatedCopilotCredits = undefined;
+		}
 		let availableTools = await this.getAvailableTools(outputStream, token);
 
 		// Emit tools_available on the agent span once, before the first CHAT span
@@ -1637,11 +1652,18 @@ export abstract class ToolCallingLoop<TOptions extends IToolCallingLoopOptions =
 		// Report token usage to the stream for rendering the context window widget
 		const stream = streamParticipants[streamParticipants.length - 1];
 		if (fetchResult.type === ChatFetchResponseType.Success && fetchResult.usage && stream && this.shouldReportUsageToContextWidget()) {
+			// Credits are billed per model call and a single turn can make many calls.
+			// Accumulate so the per-request usage reflects the whole turn (and the
+			// session cost sums correctly) instead of only the final call's credits.
+			const callCredits = nanoAiuToCredits(fetchResult.usage.copilot_usage?.total_nano_aiu);
+			if (callCredits !== undefined) {
+				this._accumulatedCopilotCredits = (this._accumulatedCopilotCredits ?? 0) + callCredits;
+			}
 			stream.usage({
 				completionTokens: fetchResult.usage.completion_tokens,
 				promptTokens: fetchResult.usage.prompt_tokens,
 				outputBuffer: endpoint.maxOutputTokens,
-				copilotCredits: nanoAiuToCredits(fetchResult.usage.copilot_usage?.total_nano_aiu),
+				copilotCredits: this._accumulatedCopilotCredits,
 				promptTokenDetails,
 			});
 		}

extensions/copilot/src/extension/intents/test/node/toolCallingLoopUsage.spec.ts

@@ -21,10 +21,10 @@ import { createExtensionUnitTestingServices } from '../../../test/node/services'
 import { IToolCallingLoopOptions, ToolCallingLoop } from '../../node/toolCallingLoop';
 
 class UsageCapturingStream extends ChatResponseStreamImpl {
-	public readonly usages: Array<{ promptTokens: number; completionTokens: number }>;
+	public readonly usages: Array<{ promptTokens: number; completionTokens: number; copilotCredits: number | undefined }>;
 
 	constructor() {
-		const usages: Array<{ promptTokens: number; completionTokens: number }> = [];
+		const usages: Array<{ promptTokens: number; completionTokens: number; copilotCredits: number | undefined }> = [];
 		super(
 			() => { },
 			() => { },
@@ -35,7 +35,8 @@ class UsageCapturingStream extends ChatResponseStreamImpl {
 			(usage) => {
 				usages.push({
 					promptTokens: usage.promptTokens,
-					completionTokens: usage.completionTokens
+					completionTokens: usage.completionTokens,
+					copilotCredits: usage.copilotCredits
 				});
 			}
 		);
@@ -71,6 +72,36 @@ class UsageTestToolCallingLoop extends ToolCallingLoop<IToolCallingLoopOptions>
 	}
 }
 
+class CreditsTestToolCallingLoop extends ToolCallingLoop<IToolCallingLoopOptions> {
+	protected override async buildPrompt(_buildPromptContext: IBuildPromptContext): Promise<IBuildPromptResult> {
+		return {
+			...nullRenderPromptResult(),
+			messages: [{ role: Raw.ChatRole.User, content: [toTextPart('hello world')] }],
+		};
+	}
+
+	protected override async getAvailableTools(): Promise<LanguageModelToolInformation[]> {
+		return [];
+	}
+
+	// Each model call bills 5 credits (5 * 1e9 nano-AIU).
+	protected override async fetch(): Promise<ChatResponse> {
+		return {
+			type: ChatFetchResponseType.Success,
+			value: 'test-response',
+			requestId: 'request-id',
+			serverRequestId: undefined,
+			usage: {
+				prompt_tokens: 100,
+				completion_tokens: 20,
+				total_tokens: 120,
+				copilot_usage: { total_nano_aiu: 5_000_000_000 }
+			},
+			resolvedModel: 'gpt-4.1'
+		};
+	}
+}
+
 const chatPanelLocation: ChatRequest['location'] = 1;
 
 function createMockChatRequest(overrides: Partial<ChatRequest> = {}): ChatRequest {
@@ -136,7 +167,7 @@ describe('ToolCallingLoop usage reporting', () => {
 
 		await loop.runOne(stream, 0, tokenSource.token);
 
-		expect(stream.usages).toEqual([{ promptTokens: 100, completionTokens: 20 }]);
+		expect(stream.usages).toEqual([{ promptTokens: 100, completionTokens: 20, copilotCredits: undefined }]);
 	});
 
 	it('does not report usage for subagent requests', async () => {
@@ -159,4 +190,25 @@ describe('ToolCallingLoop usage reporting', () => {
 
 		expect(stream.usages).toHaveLength(0);
 	});
+
+	it('accumulates copilot credits across iterations within a turn', async () => {
+		const request = createMockChatRequest();
+		const loop = instantiationService.createInstance(
+			CreditsTestToolCallingLoop,
+			{
+				conversation: createConversation(request.prompt),
+				toolCallLimit: 5,
+				request,
+			}
+		);
+		disposables.add(loop);
+		const stream = new UsageCapturingStream();
+
+		// Two model calls in the same turn: the per-request credits must be the
+		// running total (5, then 10), not just the final call's 5.
+		await loop.runOne(stream, 0, tokenSource.token);
+		await loop.runOne(stream, 1, tokenSource.token);
+
+		expect(stream.usages.map(u => u.copilotCredits)).toEqual([5, 10]);
+	});
 });