Fix context size handling for Agent Host

src/vs/platform/agentHost/common/agentService.ts

@@ -701,6 +701,8 @@ export interface IAgentModelInfo {
 	readonly id: string;
 	readonly name: string;
 	readonly maxContextWindow?: number;
+	readonly maxOutputTokens?: number;
+	readonly maxPromptTokens?: number;
 	readonly supportsVision: boolean;
 	readonly configSchema?: ConfigSchema;
 	readonly policyState?: PolicyState;

src/vs/platform/agentHost/common/state/protocol/channels-root/state.ts

@@ -6,7 +6,7 @@
 // allow-any-unicode-comment-file
 // DO NOT EDIT -- auto-generated by scripts/sync-agent-host-protocol.ts
 
-import type { ConfigSchema, ProtectedResourceMetadata } from '../common/state.js';
+import type { ConfigSchema, JsonPrimitive, ProtectedResourceMetadata } from '../common/state.js';
 import type { TerminalInfo } from '../channels-terminal/state.js';
 import type { Customization } from '../channels-session/state.js';
 
@@ -97,6 +97,10 @@ export interface SessionModelInfo {
 	name: string;
 	/** Maximum context window size */
 	maxContextWindow?: number;
+	/** Maximum number of output tokens the model can generate */
+	maxOutputTokens?: number;
+	/** Maximum number of prompt (input) tokens the model accepts */
+	maxPromptTokens?: number;
 	/** Whether the model supports vision */
 	supportsVision?: boolean;
 	/** Policy configuration state */
@@ -126,8 +130,12 @@ export interface SessionModelInfo {
 export interface ModelSelection {
 	/** Model identifier */
 	id: string;
-	/** Model-specific configuration values */
-	config?: Record<string, string>;
+	/**
+	 * Model-specific configuration values. Values are JSON primitives: most
+	 * pickers produce strings, but some (e.g. a numeric context-size picker)
+	 * produce numbers or booleans, which are carried through as-is.
+	 */
+	config?: Record<string, JsonPrimitive>;
 }
 
 // ─── Root Config Types ───────────────────────────────────────────────────────

src/vs/platform/agentHost/common/state/protocol/common/state.ts

@@ -27,6 +27,9 @@ export type URI = string;
  */
 export type StringOrMarkdown = string | { markdown: string };
 
+/** A primitive JSON value: a string, number, boolean, or `null`. */
+export type JsonPrimitive = string | number | boolean | null;
+
 // ─── Icon ────────────────────────────────────────────────────────────────────
 
 /**
@@ -161,8 +164,8 @@ export interface ConfigPropertySchema {
 	description?: string;
 	/** JSON Schema: default value */
 	default?: unknown;
-	/** JSON Schema: allowed values (typically used with `string` type) */
-	enum?: string[];
+	/** JSON Schema: allowed values. May be primitives of any JSON type. */
+	enum?: JsonPrimitive[];
 	/** Display extension: human-readable label per enum value (parallel array) */
 	enumLabels?: string[];
 	/** Display extension: description per enum value (parallel array) */

src/vs/platform/agentHost/node/agentSideEffects.ts

@@ -177,6 +177,8 @@ export class AgentSideEffects extends Disposable {
 					provider: m.provider,
 					name: m.name,
 					maxContextWindow: m.maxContextWindow,
+					maxOutputTokens: m.maxOutputTokens,
+					maxPromptTokens: m.maxPromptTokens,
 					supportsVision: m.supportsVision,
 					policyState: m.policyState,
 					configSchema: m.configSchema,

src/vs/platform/agentHost/node/claude/claudeAgent.ts

@@ -107,6 +107,8 @@ function toAgentModelInfo(m: CCAModel, provider: AgentProvider): IAgentModelInfo
 		id: m.id,
 		name: m.name,
 		maxContextWindow: m.capabilities?.limits?.max_context_window_tokens,
+		maxOutputTokens: m.capabilities?.limits?.max_output_tokens,
+		maxPromptTokens: m.capabilities?.limits?.max_prompt_tokens,
 		supportsVision: !!supports?.vision,
 		...(configSchema ? { configSchema } : {}),
 		...(policyState ? { policyState } : {}),

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

@@ -782,6 +782,8 @@ export class CodexAgent extends Disposable implements IAgent {
 					id: m.id,
 					name: m.name ?? m.id,
 					maxContextWindow: m.capabilities?.limits?.max_context_window_tokens,
+					maxOutputTokens: m.capabilities?.limits?.max_output_tokens,
+					maxPromptTokens: m.capabilities?.limits?.max_prompt_tokens,
 					supportsVision: !!m.capabilities?.supports?.vision,
 					configSchema,
 					policyState: m.policy?.state as PolicyState | undefined,

src/vs/platform/agentHost/node/copilot/copilotAgent.ts

@@ -53,7 +53,7 @@ import { ICopilotBranchNameGenerator } from './copilotBranchNameGenerator.js';
 import { CopilotAgentSession, type CopilotSdkMode } from './copilotAgentSession.js';
 import { ICopilotSessionContext, projectFromCopilotContext } from './copilotGitProject.js';
 import { parsedPluginsEqual, toChildCustomizations } from './copilotPluginConverters.js';
-import { CopilotSessionLauncher, ContextTierConfigKey, ThinkingLevelConfigKey, getCopilotContextTier, getCopilotReasoningEffort, type CopilotSessionLaunchPlan, type IActiveClientSnapshot } from './copilotSessionLauncher.js';
+import { CopilotSessionLauncher, ContextSizeConfigKey, ThinkingLevelConfigKey, getCopilotContextTier, getCopilotReasoningEffort, type CopilotSessionLaunchPlan, type IActiveClientSnapshot } from './copilotSessionLauncher.js';
 import { ShellManager } from './copilotShellTools.js';
 import { isRestrictedTelemetryEnabled } from './copilotTokenFields.js';
 import { CopilotSlashCommandCompletionProvider } from './copilotSlashCommandCompletionProvider.js';
@@ -809,14 +809,19 @@ export class CopilotAgent extends Disposable implements IAgent {
 	}
 
 	/**
-	 * Synthesize a `contextTier` config property when the model exposes a `long_context` pricing tier with a distinct
+	 * Synthesize a `contextSize` config property when the model exposes a `long_context` pricing tier with a distinct
 	 * context-max. Picker surfaces this as the "Context Size" button. Mirrors `getContextSizeOptions` in
 	 * `extensions/copilot/src/extension/conversation/vscode-node/languageModelAccess.ts`.
 	 *
+	 * The `enum` values are the two context-window sizes (in tokens), smallest first, so the numeric token counts
+	 * flow to the client. The chosen value comes back in the model's `config` bag and is mapped to the SDK's
+	 * two-valued `contextTier` at the SDK boundary by {@link getCopilotContextTier}, using the model's long-context
+	 * window from {@link _longContextWindowFor}.
+	 *
 	 * `billing.tokenPrices` is present on the runtime CAPI `/models` payload but not yet declared on the published SDK
 	 * `ModelBilling` type — narrow through {@link ICAPIModelBilling} until the SDK catches up.
 	 */
-	private _createContextTierConfigSchemaProperty(billing: ModelInfo['billing'] | undefined): ConfigPropertySchema | undefined {
+	private _createContextSizeConfigSchemaProperty(billing: ModelInfo['billing'] | undefined): ConfigPropertySchema | undefined {
 		const tokenPrices = billing?.tokenPrices;
 		const defaultMax = tokenPrices?.contextMax;
 		const longContextMax = tokenPrices?.longContext?.contextMax;
@@ -828,21 +833,37 @@ export class CopilotAgent extends Disposable implements IAgent {
 			|| typeof tokenPrices?.longContext?.outputPrice === 'number';
 
 		return {
-			type: 'string',
-			title: localize('copilot.modelContextTier.title', "Context Size"),
-			description: localize('copilot.modelContextTier.description', "Selects the context window size for this model."),
-			default: 'default',
-			enum: ['default', 'long_context'],
+			type: 'number',
+			title: localize('copilot.modelContextSize.title', "Context Size"),
+			description: localize('copilot.modelContextSize.description', "Selects the context window size for this model."),
+			default: defaultMax,
+			enum: [defaultMax, longContextMax],
 			enumLabels: [formatTokenCount(defaultMax), formatTokenCount(longContextMax)],
 			enumDescriptions: [
-				localize('copilot.modelContextTier.default', "Default"),
+				localize('copilot.modelContextSize.default', "Default"),
 				hasLongContextSurcharge
-					? localize('copilot.modelContextTier.longerSessions', "Longer sessions")
-					: localize('copilot.modelContextTier.longerSessionsNoCompaction', "Longer sessions without compaction"),
+					? localize('copilot.modelContextSize.longerSessions', "Longer sessions")
+					: localize('copilot.modelContextSize.longerSessionsNoCompaction', "Longer sessions without compaction"),
 			],
 		};
 	}
 
+	/**
+	 * The model's long-context window (in tokens): the largest size offered by its "Context Size" picker
+	 * (the max numeric value in the synthesized `contextSize` {@link ConfigPropertySchema.enum}). Used by
+	 * {@link getCopilotContextTier} to decide whether a numeric selection opts into `long_context`.
+	 * Returns `undefined` when the model exposes no such picker (or the model list isn't loaded yet),
+	 * leaving the SDK on its default tier.
+	 */
+	private _longContextWindowFor(modelId: string | undefined): number | undefined {
+		if (!modelId) {
+			return undefined;
+		}
+		const windows = this._models.get().find(m => m.id === modelId)?.configSchema?.properties?.[ContextSizeConfigKey]?.enum;
+		const numericWindows = windows?.filter((w): w is number => typeof w === 'number');
+		return numericWindows && numericWindows.length > 0 ? Math.max(...numericWindows) : undefined;
+	}
+
 	/**
 	 * Builds the open `_meta` pricing bag for a model from its billing info so the chat model picker can render its
 	 * cost hover. Delegates to the shared {@link createPricingMetaFromBilling} helper.
@@ -859,9 +880,9 @@ export class CopilotAgent extends Disposable implements IAgent {
 		if (thinkingLevel) {
 			properties[ThinkingLevelConfigKey] = thinkingLevel;
 		}
-		const contextTier = this._createContextTierConfigSchemaProperty(m.billing);
-		if (contextTier) {
-			properties[ContextTierConfigKey] = contextTier;
+		const contextSize = this._createContextSizeConfigSchemaProperty(m.billing);
+		if (contextSize) {
+			properties[ContextSizeConfigKey] = contextSize;
 		}
 		if (Object.keys(properties).length === 0) {
 			return undefined;
@@ -1023,6 +1044,8 @@ export class CopilotAgent extends Disposable implements IAgent {
 			// Synthetic SDK entries like `auto` ship with `capabilities: {}` and
 			// no fixed context window — surface them with maxContextWindow undefined.
 			maxContextWindow: m.capabilities?.limits?.max_context_window_tokens,
+			maxOutputTokens: m.capabilities?.limits?.max_output_tokens,
+			maxPromptTokens: m.capabilities?.limits?.max_prompt_tokens,
 			supportsVision: !!m.capabilities?.supports?.vision,
 			configSchema: this._createModelConfigSchema(m),
 			policyState: m.policy?.state as PolicyState | undefined,
@@ -1252,6 +1275,7 @@ export class CopilotAgent extends Disposable implements IAgent {
 				shellManager,
 				githubToken: this._githubToken,
 				model: provisional.model,
+				longContextWindow: this._longContextWindowFor(provisional.model?.id),
 			};
 			agentSession = this._createAgentSession(launchPlan, customizationDirectory, activeClient);
 			await agentSession.initializeSession();
@@ -1719,6 +1743,7 @@ export class CopilotAgent extends Disposable implements IAgent {
 			if (this._chatSessions.has(chatKey)) {
 				return;
 			}
+			const model = options?.model;
 			// Resolve the owning session so the new chat inherits its working
 			// directory scope. The parent may be provisional (no SDK session
 			// yet); in that case use its provisional working directory.
@@ -1762,7 +1787,7 @@ export class CopilotAgent extends Disposable implements IAgent {
 					activeClientToolSet: activeClient.toolSet,
 					shellManager,
 					githubToken: this._githubToken,
-					fallback: { model: options.model },
+					fallback: { model, longContextWindow: this._longContextWindowFor(model?.id) },
 				};
 			} else {
 				sdkSessionId = chatSdkId;
@@ -1776,7 +1801,8 @@ export class CopilotAgent extends Disposable implements IAgent {
 					activeClientToolSet: activeClient.toolSet,
 					shellManager,
 					githubToken: this._githubToken,
-					model: options?.model,
+					model,
+					longContextWindow: this._longContextWindowFor(model?.id),
 				};
 			}
 			let agentSession: CopilotAgentSession | undefined;
@@ -1790,7 +1816,7 @@ export class CopilotAgent extends Disposable implements IAgent {
 				const parsed = parseChatUri(chat);
 				if (parsed) {
 					const persisted = await this._readPersistedChats(session);
-					persisted.set(parsed.chatId, { sdkSessionId, ...(options?.model ? { model: options.model } : {}) });
+					persisted.set(parsed.chatId, { sdkSessionId, ...(model ? { model } : {}) });
 					await this._writePersistedChats(session, persisted);
 				}
 				this._logService.info(`[Copilot] Created additional chat ${chatKey} in session ${session.toString()}${options?.fork ? ' (forked)' : ''}`);
@@ -1941,7 +1967,7 @@ export class CopilotAgent extends Disposable implements IAgent {
 				activeClientToolSet: activeClient.toolSet,
 				shellManager,
 				githubToken: this._githubToken,
-				fallback: { model: info.model },
+				fallback: { model: info.model, longContextWindow: this._longContextWindowFor(info.model?.id) },
 			};
 			let agentSession: CopilotAgentSession | undefined;
 			try {
@@ -1992,12 +2018,13 @@ export class CopilotAgent extends Disposable implements IAgent {
 	}
 
 	async changeModel(session: URI, model: ModelSelection, chat?: URI): Promise<void> {
+		const longContextWindow = this._longContextWindowFor(model.id);
 		// Additional (non-default) chats are backed by their own SDK
 		// conversation tracked in `_chatSessions`; apply the change there and
 		// skip the session-level metadata store (peer chats are not persisted
 		// per-chat).
 		if (chat && !isDefaultChatUri(chat)) {
-			await this._chatSessions.get(chat.toString())?.setModel(model.id, getCopilotReasoningEffort(model), getCopilotContextTier(model));
+			await this._chatSessions.get(chat.toString())?.setModel(model.id, getCopilotReasoningEffort(model), getCopilotContextTier(model, longContextWindow));
 			return;
 		}
 		const sessionId = AgentSession.id(session);
@@ -2008,7 +2035,7 @@ export class CopilotAgent extends Disposable implements IAgent {
 		}
 		const entry = this._sessions.get(sessionId);
 		if (entry) {
-			await entry.setModel(model.id, getCopilotReasoningEffort(model), getCopilotContextTier(model));
+			await entry.setModel(model.id, getCopilotReasoningEffort(model), getCopilotContextTier(model, longContextWindow));
 		}
 		await this._storeSessionMetadata(session, model, undefined, undefined, undefined);
 	}
@@ -2264,6 +2291,7 @@ export class CopilotAgent extends Disposable implements IAgent {
 			githubToken: this._githubToken,
 			fallback: {
 				model: storedMetadata.model,
+				longContextWindow: this._longContextWindowFor(storedMetadata.model?.id),
 			},
 		};
 

src/vs/platform/agentHost/node/copilot/copilotSessionLauncher.ts

@@ -27,6 +27,16 @@ import { describeSystemMessageConfig } from './prompts/systemMessage.js';
 import './prompts/allPrompts.js';
 
 export const ThinkingLevelConfigKey = 'thinkingLevel';
+/**
+ * Config key for the numeric "Context Size" selection (a context-window token count). Mapped to the
+ * SDK's two-valued {@link SessionConfig.contextTier} by {@link getCopilotContextTier}.
+ */
+export const ContextSizeConfigKey = 'contextSize';
+/**
+ * @deprecated Legacy config key that stored the resolved tier string (`'default'` / `'long_context'`)
+ * directly. Replaced by the numeric {@link ContextSizeConfigKey}; still read from persisted sessions
+ * for backward compatibility.
+ */
 export const ContextTierConfigKey = 'contextTier';
 
 const ReasoningEfforts = ['low', 'medium', 'high', 'xhigh'] as const;
@@ -120,23 +130,25 @@ interface ICopilotSessionLaunchBase {
 export interface ICopilotCreateSessionLaunchPlan extends ICopilotSessionLaunchBase {
 	readonly kind: 'create';
 	readonly model: ModelSelection | undefined;
+	readonly longContextWindow?: number;
 }
 
 export interface ICopilotResumeSessionLaunchPlan extends ICopilotSessionLaunchBase {
 	readonly kind: 'resume';
 	readonly workingDirectory: URI;
 	readonly fallback: {
 		readonly model: ModelSelection | undefined;
+		readonly longContextWindow?: number;
 	};
 }
 
 export type CopilotSessionLaunchPlan = ICopilotCreateSessionLaunchPlan | ICopilotResumeSessionLaunchPlan;
 
-function isReasoningEffort(value: string | undefined): value is ReasoningEffort {
+function isReasoningEffort(value: unknown): value is ReasoningEffort {
 	return ReasoningEfforts.some(reasoningEffort => reasoningEffort === value);
 }
 
-function isContextTier(value: string | undefined): value is ContextTier {
+function isContextTier(value: unknown): value is ContextTier {
 	return ContextTiers.some(contextTier => contextTier === value);
 }
 
@@ -189,9 +201,26 @@ export function getCopilotReasoningEffort(model: ModelSelection | undefined): Se
 	return isReasoningEffort(thinkingLevel) ? thinkingLevel : undefined;
 }
 
-export function getCopilotContextTier(model: ModelSelection | undefined): SessionConfig['contextTier'] {
-	const contextTier = model?.config?.[ContextTierConfigKey];
-	return isContextTier(contextTier) ? contextTier : undefined;
+export function getCopilotContextTier(model: ModelSelection | undefined, longContextWindow?: number): SessionConfig['contextTier'] {
+	// Legacy persisted selections stored the resolved tier string directly under the deprecated key.
+	const legacyTier = model?.config?.[ContextTierConfigKey];
+	if (isContextTier(legacyTier)) {
+		return legacyTier;
+	}
+	// The "Context Size" picker exposes numeric token-count enum values, so a current selection arrives
+	// under `contextSize` as a token count. Map it to the SDK's two-valued tier using the model's
+	// long-context window: only a selection that reaches that window opts into `long_context`. Without
+	// the window (model exposes no picker, or the model list isn't loaded) leave the SDK on its default
+	// tier.
+	const contextSize = model?.config?.[ContextSizeConfigKey];
+	if (contextSize === undefined) {
+		return undefined;
+	}
+	const selectedWindow = Number(contextSize);
+	if (!Number.isFinite(selectedWindow) || typeof longContextWindow !== 'number') {
+		return undefined;
+	}
+	return selectedWindow >= longContextWindow ? 'long_context' : 'default';
 }
 
 export class CopilotSessionLauncher implements ICopilotSessionLauncher {
@@ -236,6 +265,7 @@ export class CopilotSessionLauncher implements ICopilotSessionLauncher {
 				...plan,
 				kind: 'create',
 				model: plan.fallback.model,
+				longContextWindow: plan.fallback.longContextWindow,
 			}, config, sandboxConfig);
 			this._logService.info(`[Copilot:${plan.sessionId}] Fallback createSession succeeded`);
 			return wrapper;
@@ -249,7 +279,7 @@ export class CopilotSessionLauncher implements ICopilotSessionLauncher {
 			streaming: true,
 			model: plan.model?.id,
 			reasoningEffort: getCopilotReasoningEffort(plan.model),
-			contextTier: getCopilotContextTier(plan.model),
+			contextTier: getCopilotContextTier(plan.model, plan.longContextWindow),
 			...(plan.resolvedAgentName ? { agent: plan.resolvedAgentName } : {}),
 			workingDirectory: plan.workingDirectory?.fsPath,
 		});

src/vs/platform/agentHost/test/node/agentSideEffects.test.ts

@@ -707,7 +707,7 @@ suite('AgentSideEffects', () => {
 				}
 				return e.action.agents[0]?.models.length === 1;
 			}));
-			agent.setModels([{ provider: 'mock', id: 'mock-model', name: 'mock Model', maxContextWindow: 128000, supportsVision: false }]);
+			agent.setModels([{ provider: 'mock', id: 'mock-model', name: 'mock Model', maxContextWindow: 128000, maxOutputTokens: 16000, maxPromptTokens: 112000, supportsVision: false }]);
 			await envelope;
 
 			const actions = envelopes.map(e => e.action).filter(action => action.type === ActionType.RootAgentsChanged);
@@ -718,6 +718,8 @@ suite('AgentSideEffects', () => {
 				provider: 'mock',
 				name: 'mock Model',
 				maxContextWindow: 128000,
+				maxOutputTokens: 16000,
+				maxPromptTokens: 112000,
 				supportsVision: false,
 				policyState: undefined,
 				configSchema: undefined,