agentHost: fix context picker for ChatGPT Codex models

src/vs/platform/agentHost/node/codex/codexAgent.ts

@@ -26,7 +26,7 @@ import { IProductService } from '../../../product/common/productService.js';
 import { createSchema, platformRootSchema, platformSessionSchema, schemaProperty, AgentHostAutoApprovePolicyRestrictedConfigKey, AgentHostCodexMultiRootEnabledConfigKey, AgentHostGitHubMcpServerEnabledConfigKey, AgentHostMcpServersConfigKey, type ISchemaProperty, type SessionMode } from '../../common/agentHostSchema.js';
 import { createPricingMetaFromBilling, normalizeCAPIBilling, type ICAPIModelBilling } from '../../common/agentModelPricing.js';
 import { ContextSizeConfigKey, createContextSizeConfigSchemaProperty, getModelContextSize } from '../../common/agentModelConfiguration.js';
-import { CHATGPT_SUBSCRIPTION_MODEL_SOURCE_ID, createAgentModelGroupMeta, createAgentModelSourceMeta } from '../../common/agentModelSource.js';
+import { CHATGPT_SUBSCRIPTION_MODEL_SOURCE_ID, createAgentModelGroupMeta, createAgentModelSourceMeta, readAgentModelSourceId } from '../../common/agentModelSource.js';
 import { AgentSystemNotificationKind, toAgentSystemNotificationMeta } from '../../common/meta/agentSystemNotificationMeta.js';
 import { AgentHostConfigKey, agentHostCustomizationConfigSchema } from '../../common/agentHostCustomizationConfig.js';
 import { AgentSdkSetupChannel } from '../agentSdkSetupChannel.js';
@@ -1760,6 +1760,28 @@ export class CodexAgent extends Disposable implements IAgent {
 		);
 	}
 
+	private _withCopilotContextSize(model: IAgentModelInfo): IAgentModelInfo {
+		if (readAgentModelSourceId(model) !== CHATGPT_SUBSCRIPTION_MODEL_SOURCE_ID) {
+			return model;
+		}
+		const modelId = parseCodexModelSelection(model).modelId;
+		const copilotModel = this._copilotModels.find(candidate => parseCodexModelSelection(candidate).modelId === modelId);
+		const contextSize = copilotModel?.configSchema?.properties[ContextSizeConfigKey];
+		if (!contextSize) {
+			return model;
+		}
+		return {
+			...model,
+			configSchema: {
+				type: 'object',
+				properties: {
+					...model.configSchema?.properties,
+					[ContextSizeConfigKey]: contextSize,
+				},
+			},
+		};
+	}
+
 	/**
 	 * Resolve the Codex security axes (approval policy, sandbox, reviewer) for a
 	 * live or restored session from its RAW persisted config values.
@@ -1968,7 +1990,7 @@ export class CodexAgent extends Disposable implements IAgent {
 		if (generation !== this._modelCatalogGeneration || this._isShuttingDown || this._store.isDisposed) {
 			return;
 		}
-		this._models.set([...this._copilotModels, ...this._codexModels], undefined);
+		this._models.set([...this._copilotModels, ...this._codexModels.map(model => this._withCopilotContextSize(model))], undefined);
 		// Last, never first: announcing `ready` before the catalog lands is how the
 		// window renders "no account found".
 		this._sdkSetupChannel.refresh();

src/vs/platform/agentHost/test/node/codex/codexModelRefresh.test.ts

@@ -1527,6 +1527,71 @@ suite('CodexAgent model refresh', () => {
 		}]);
 	});
 
+	test('publishes context size options for ChatGPT subscription models', async () => {
+		const copilotModel: CCAModel = {
+			billing: {
+				is_premium: true,
+				multiplier: 1,
+				restricted_to: [],
+				token_prices: {
+					default: { context_max: 272_000, input_price: 1 },
+					long_context: { context_max: 1_000_000, input_price: 2 },
+				},
+			},
+			capabilities: {
+				family: 'gpt-5.6',
+				limits: { max_context_window_tokens: 272_000, max_output_tokens: 32_000, max_prompt_tokens: 240_000 },
+				object: 'model_capabilities',
+				supports: { parallel_tool_calls: true, streaming: true, tool_calls: true, vision: true },
+				tokenizer: 'o200k_base',
+				type: 'chat',
+			},
+			id: 'gpt-5.6-sol',
+			is_chat_default: true,
+			is_chat_fallback: false,
+			model_picker_category: 'advanced',
+			name: 'GPT-5.6-Sol',
+			model_picker_enabled: true,
+			object: 'model',
+			policy: { state: 'enabled', terms: '' },
+			preview: false,
+			supported_endpoints: ['/responses'],
+			vendor: 'OpenAI',
+			version: 'gpt-5.6-sol',
+		};
+		let copilotModels = [copilotModel];
+		const agent = createAgent(disposables, async () => copilotModels);
+		agent['_githubToken'] = 'token';
+		agent['_connection'] = createChatGPTConnection() as never;
+
+		await agent.refreshModels();
+
+		assert.deepStrictEqual(agent.models.get().map(model => ({
+			id: model.id,
+			contextSize: model.configSchema?.properties.contextSize && {
+				enum: model.configSchema.properties.contextSize.enum,
+				default: model.configSchema.properties.contextSize.default,
+			},
+		})), [{
+			id: toCodexModelSelectionId('vscode-proxy', 'gpt-5.6-sol'),
+			contextSize: {
+				enum: [272_000, 1_000_000],
+				default: 272_000,
+			},
+		}, {
+			id: toCodexModelSelectionId('openai', 'gpt-5.6-sol'),
+			contextSize: {
+				enum: [272_000, 1_000_000],
+				default: 272_000,
+			},
+		}]);
+
+		copilotModels = [{ ...copilotModel, id: 'gpt-5.6-terra' }];
+		await agent.refreshModels();
+		const chatGPTModel = agent.models.get().find(model => model.id === toCodexModelSelectionId('openai', 'gpt-5.6-sol'));
+		assert.strictEqual(chatGPTModel?.configSchema?.properties.contextSize, undefined);
+	});
+
 	test('omits the thinking level when a Codex model advertises no reasoning efforts', async () => {
 		const agent = createAgent(disposables, async () => []);
 		agent['_connection'] = {

src/vs/platform/agentHost/test/node/codex/codexPrewarmEviction.test.ts

@@ -1924,68 +1924,95 @@ suite('CodexAgent prewarm eviction', () => {
 		peer.exit();
 	});
 
-	test('applies context size and thinking level to new and resumed threads', async () => {
-		const agent = await createAgent(disposables);
-		agent['_schedulePrewarm'] = () => { };
-		agent['_refreshSkillHookCustomizations'] = async () => { };
-		agent['_refreshSkillExtraRoots'] = async () => { };
-		const peer = disposables.add(createTestPeer());
-		agent['_connection'] = {
-			kind: 'ready',
-			client: new CodexAppServerClient(peer.transport),
-			usageSource: 'github',
-			child: { kill: () => true },
-		} as never;
-		const baseModel = agent.models.get()[0];
-		agent['_models'].set([{
-			...baseModel,
-			configSchema: {
-				type: 'object',
-				properties: {
-					thinkingLevel: { type: 'string', title: 'Thinking Level', enum: ['low', 'high'], default: 'low' },
-					contextSize: { type: 'number', title: 'Context Size', enum: [272_000, 1_000_000], default: 272_000 },
+	test('applies context size and thinking level for Copilot and ChatGPT models', async () => {
+		const runScenario = async (source: 'copilot' | 'chatgpt', selectedModelId: string) => {
+			const agent = await createAgent(disposables);
+			agent['_schedulePrewarm'] = () => { };
+			agent['_refreshSkillHookCustomizations'] = async () => { };
+			agent['_refreshSkillExtraRoots'] = async () => { };
+			const peer = disposables.add(createTestPeer());
+			agent['_connection'] = {
+				kind: 'ready',
+				client: new CodexAppServerClient(peer.transport),
+				usageSource: source === 'copilot' ? 'github' : 'openai',
+				child: { kill: () => true },
+			} as never;
+			const baseModel = agent.models.get()[0];
+			agent['_models'].set([{
+				...baseModel,
+				id: selectedModelId,
+				configSchema: {
+					type: 'object',
+					properties: {
+						thinkingLevel: { type: 'string', title: 'Thinking Level', enum: ['low', 'high'], default: 'low' },
+						contextSize: { type: 'number', title: 'Context Size', enum: [272_000, 1_000_000], default: 272_000 },
+					},
 				},
-			},
-		}], undefined);
+			}], undefined);
 
-		const folder = URI.file('/repo/context-size');
-		const longContextModel = { id: COPILOT_TEST_MODEL, config: { thinkingLevel: 'low', contextSize: 1_000_000 } };
-		const created = await createSession(agent, { workingDirectories: [folder], model: longContextModel });
-		const chat = defaultChatOf(created.session);
-		const entry = agent['_sessions'].get(AgentSession.id(created.session))!;
-		const materializing = agent['_materializeIfNeeded'](entry, created.session, false);
-		const start = await readNextRequest(peer.outbound);
-		peer.push({ id: start.id, result: { thread: { id: 'context-size-thread', cwd: folder.fsPath } } });
-		await materializing;
+			const threadId = `${source}-context-size-thread`;
+			const folder = URI.file(`/repo/context-size-${source}`);
+			const longContextModel = { id: selectedModelId, config: { thinkingLevel: 'low', contextSize: 1_000_000 } };
+			const created = await createSession(agent, { workingDirectories: [folder], model: longContextModel });
+			const chat = defaultChatOf(created.session);
+			const entry = agent['_sessions'].get(AgentSession.id(created.session))!;
+			const materializing = agent['_materializeIfNeeded'](entry, created.session, false);
+			const start = await readNextRequest(peer.outbound);
+			peer.push({ id: start.id, result: { thread: { id: threadId, cwd: folder.fsPath } } });
+			await materializing;
 
-		await agent.chats.changeModel(chat, { id: COPILOT_TEST_MODEL, config: { thinkingLevel: 'high', contextSize: 272_000 } }, chatContext(created.session, chat));
-		const sending = agent.chats.sendMessage(chat, 'use the shorter window', [folder], undefined, 'turn-1', undefined, undefined, chatContext(created.session, chat));
-		const unsubscribe = await readNextRequest(peer.outbound);
-		peer.push({ id: unsubscribe.id, result: {} });
-		const resume = await readNextRequest(peer.outbound);
-		peer.push({ id: resume.id, result: { thread: { id: 'context-size-thread', cwd: folder.fsPath }, cwd: folder.fsPath } });
-		const inventory = await readNextRequest(peer.outbound);
-		peer.push({ id: inventory.id, result: { data: [], nextCursor: null } });
-		const turn = await readNextRequest(peer.outbound);
-		peer.push({ id: turn.id, result: {} });
-		await sending;
+			await agent.chats.changeModel(chat, { id: selectedModelId, config: { thinkingLevel: 'high', contextSize: 272_000 } }, chatContext(created.session, chat));
+			const sending = agent.chats.sendMessage(chat, 'use the shorter window', [folder], undefined, 'turn-1', undefined, undefined, chatContext(created.session, chat));
+			const unsubscribe = await readNextRequest(peer.outbound);
+			peer.push({ id: unsubscribe.id, result: {} });
+			const resume = await readNextRequest(peer.outbound);
+			peer.push({ id: resume.id, result: { thread: { id: threadId, cwd: folder.fsPath }, cwd: folder.fsPath } });
+			const inventory = await readNextRequest(peer.outbound);
+			peer.push({ id: inventory.id, result: { data: [], nextCursor: null } });
+			const turn = await readNextRequest(peer.outbound);
+			peer.push({ id: turn.id, result: {} });
+			await sending;
+			peer.exit();
 
-		assert.deepStrictEqual({
-			start: { method: start.method, contextSize: start.params.config?.model_context_window },
-			unsubscribe: { method: unsubscribe.method, threadId: unsubscribe.params.threadId },
-			resume: { method: resume.method, contextSize: resume.params.config?.model_context_window },
-			turn: {
-				method: turn.method,
-				thinkingLevel: turn.params.effort,
-				collaborationThinkingLevel: turn.params.collaborationMode?.settings.reasoning_effort,
-			},
+			return {
+				source,
+				start: {
+					method: start.method,
+					model: start.params.model,
+					modelProvider: start.params.modelProvider,
+					contextSize: start.params.config?.model_context_window,
+				},
+				unsubscribe: { method: unsubscribe.method, threadId: unsubscribe.params.threadId },
+				resume: {
+					method: resume.method,
+					model: resume.params.model,
+					modelProvider: resume.params.modelProvider,
+					contextSize: resume.params.config?.model_context_window,
+				},
+				turn: {
+					method: turn.method,
+					thinkingLevel: turn.params.effort,
+					collaborationThinkingLevel: turn.params.collaborationMode?.settings.reasoning_effort,
+				},
+			};
+		};
+
+		assert.deepStrictEqual([
+			await runScenario('copilot', COPILOT_TEST_MODEL),
+			await runScenario('chatgpt', OPENAI_TEST_MODEL),
+		], [{
+			source: 'copilot',
+			start: { method: 'thread/start', model: 'gpt-test', modelProvider: 'vscode-proxy', contextSize: 1_000_000 },
+			unsubscribe: { method: 'thread/unsubscribe', threadId: 'copilot-context-size-thread' },
+			resume: { method: 'thread/resume', model: 'gpt-test', modelProvider: 'vscode-proxy', contextSize: 272_000 },
+			turn: { method: 'turn/start', thinkingLevel: 'high', collaborationThinkingLevel: 'high' },
 		}, {
-			start: { method: 'thread/start', contextSize: 1_000_000 },
-			unsubscribe: { method: 'thread/unsubscribe', threadId: 'context-size-thread' },
-			resume: { method: 'thread/resume', contextSize: 272_000 },
+			source: 'chatgpt',
+			start: { method: 'thread/start', model: 'gpt-5.6-sol', modelProvider: 'openai', contextSize: 1_000_000 },
+			unsubscribe: { method: 'thread/unsubscribe', threadId: 'chatgpt-context-size-thread' },
+			resume: { method: 'thread/resume', model: 'gpt-5.6-sol', modelProvider: 'openai', contextSize: 272_000 },
 			turn: { method: 'turn/start', thinkingLevel: 'high', collaborationThinkingLevel: 'high' },
-		});
-		peer.exit();
+		}]);
 	});
 
 	test('routes provider-qualified models independently and switches one session', async () => {

src/vs/workbench/contrib/chat/browser/widget/input/chatInputModelSelectionController.ts

@@ -289,23 +289,31 @@ export class ChatInputModelSelectionController extends Disposable {
 	}
 
 	reconcileModelListChange(models: readonly ILanguageModelChatMetadataAndIdentifier[]): void {
+		const currentModel = this._currentModel.get();
+		const republishedCurrentModel = currentModel && models.find(model => model.identifier === currentModel.identifier);
+		if (republishedCurrentModel && republishedCurrentModel !== currentModel) {
+			// A provider can enrich a model after it was selected (for example, when a
+			// second catalogue supplies its context-size schema). Refresh the displayed
+			// snapshot without reapplying or persisting a selection that did not change.
+			this._display(republishedCurrentModel);
+		}
 		if (this.applyConfiguredDefault() || this._reconcilePendingProgrammaticSelection() || this._restoreRememberedModel()) {
 			return;
 		}
-		const currentModel = this._currentModel.get();
+		const reconciledCurrentModel = this._currentModel.get();
 		const declaredDefault = this._runtime.getDeclaredDefaultModel(models);
 		if (this._runtime.isEmpty()
 			&& this._selectionReason === ModelSelectionReason.FirstAvailable
 			&& declaredDefault
-			&& currentModel?.identifier !== declaredDefault.identifier) {
+			&& reconciledCurrentModel?.identifier !== declaredDefault.identifier) {
 			// Still the first thing on offer, only now the pool has said which that is.
 			this._applyModel(declaredDefault, ModelSelectionReason.FirstAvailable);
 			return;
 		}
-		if (!shouldResetOnModelListChange(currentModel?.identifier, [...models])) {
+		if (!shouldResetOnModelListChange(reconciledCurrentModel?.identifier, [...models])) {
 			return;
 		}
-		const match = findBestMatchingModel(currentModel, models);
+		const match = findBestMatchingModel(reconciledCurrentModel, models);
 		if (match) {
 			// The same selection republished under another identifier, so whoever chose it still has.
 			this._applyModel(match, this._selectionReason);

src/vs/workbench/contrib/chat/test/browser/widget/input/chatInputModelSelectionController.test.ts

@@ -202,6 +202,38 @@ suite('ChatInputModelSelectionController', () => {
 		});
 	});
 
+	test('refreshes selected model metadata when the same identifier is republished', () => {
+		const modelChanges = disposables.add(new Emitter<string>());
+		const initial = model('agent-host-codex:openai/gpt-5.6-sol');
+		const enriched = {
+			...initial,
+			metadata: {
+				...initial.metadata,
+				configurationSchema: {
+					properties: {
+						thinkingLevel: { group: 'navigation', enum: ['low', 'high'], default: 'low' },
+						contextSize: { group: 'tokens', enum: [200_000, 922_000], default: 200_000 },
+					},
+				},
+			},
+		} satisfies ILanguageModelChatMetadataAndIdentifier;
+		const state: IRuntimeState = { models: [initial], sessionType: 'agent-host-codex' };
+		const applied: string[] = [];
+		const controller = disposables.add(new ChatInputModelSelectionController(createRuntime(state, modelChanges, applied)));
+
+		controller.applySelection(initial, () => { }, false);
+		state.models = [enriched];
+		modelChanges.fire('agent-host-codex');
+
+		assert.deepStrictEqual({
+			selectedModel: controller.currentModel.get(),
+			applied,
+		}, {
+			selectedModel: enriched,
+			applied: [],
+		});
+	});
+
 	test('rolls back a failed explicit selection effect', () => {
 		const modelChanges = disposables.add(new Emitter<string>());
 		const controller = disposables.add(new ChatInputModelSelectionController(createRuntime({ models: [], sessionType: 'test' }, modelChanges, [])));