Gemini: Fix tool-call validation for Gemini flattened argument keys (#322165)

* Gemini: Fix tool-call validation for Gemini flattened argument keys

* Feedback updates

extensions/copilot/src/extension/tools/common/test/toolService.spec.ts

@@ -230,5 +230,183 @@ describe('Tool Service', () => {
 				error: expect.stringContaining('ERROR: Your input to the tool was invalid')
 			});
 		});
+
+		test('should reconstruct flattened path keys', () => {
+			const askQuestionsTool: vscode.LanguageModelToolInformation = {
+				name: 'askQuestionsTool',
+				description: 'A tool that expects an array of nested question objects',
+				inputSchema: {
+					type: 'object',
+					properties: {
+						questions: {
+							type: 'array',
+							items: {
+								type: 'object',
+								properties: {
+									header: { type: 'string' },
+									question: { type: 'string' },
+									allowFreeformInput: { type: 'boolean' },
+									options: {
+										type: 'array',
+										items: {
+											type: 'object',
+											properties: {
+												label: { type: 'string' },
+												description: { type: 'string' },
+												recommended: { type: 'boolean' }
+											},
+											required: ['label']
+										}
+									}
+								},
+								required: ['header', 'question']
+							}
+						}
+					},
+					required: ['questions']
+				},
+				tags: [],
+				source: undefined
+			};
+
+			(toolsService.tools as vscode.LanguageModelToolInformation[]).push(askQuestionsTool);
+
+			// Gemini-style flattened path keys instead of a nested object/array.
+			const flattenedInput = JSON.stringify({
+				'questions[0].allowFreeformInput': true,
+				'questions[0].header': 'repro_question_1',
+				'questions[0].options[0].description': 'First option description',
+				'questions[0].options[0].label': 'Option A',
+				'questions[0].options[0].recommended': true,
+				'questions[0].options[1].description': 'Second option description',
+				'questions[0].options[1].label': 'Option B',
+				'questions[0].question': 'Which option do you prefer?',
+				'questions[1].allowFreeformInput': false,
+				'questions[1].header': 'repro_question_2',
+				'questions[1].options[0].label': 'Yes',
+				'questions[1].options[1].label': 'No',
+				'questions[1].question': 'Do you want to continue?'
+			});
+
+			const result = toolsService.validateToolInput('askQuestionsTool', flattenedInput);
+			expect(result).toEqual({
+				inputObj: {
+					questions: [
+						{
+							allowFreeformInput: true,
+							header: 'repro_question_1',
+							question: 'Which option do you prefer?',
+							options: [
+								{ description: 'First option description', label: 'Option A', recommended: true },
+								{ description: 'Second option description', label: 'Option B' }
+							]
+						},
+						{
+							allowFreeformInput: false,
+							header: 'repro_question_2',
+							question: 'Do you want to continue?',
+							options: [
+								{ label: 'Yes' },
+								{ label: 'No' }
+							]
+						}
+					]
+				}
+			});
+		});
+
+		test('should not pollute prototype when reconstructing flattened keys', () => {
+			const pollutionTool: vscode.LanguageModelToolInformation = {
+				name: 'pollutionTool',
+				description: 'A tool whose flattened input contains unsafe property names',
+				inputSchema: {
+					type: 'object',
+					properties: {
+						data: {
+							type: 'object',
+							properties: { value: { type: 'string' } }
+						}
+					},
+					required: ['data']
+				},
+				tags: [],
+				source: undefined
+			};
+
+			(toolsService.tools as vscode.LanguageModelToolInformation[]).push(pollutionTool);
+
+			const malicious = JSON.stringify({
+				'__proto__.polluted': 'yes',
+				'data.value': 'ok'
+			});
+
+			const result = toolsService.validateToolInput('pollutionTool', malicious);
+
+			// The unsafe key makes reconstruction bail out, so validation fails
+			// rather than mutating Object.prototype.
+			expect(result).toMatchObject({
+				error: expect.stringContaining('ERROR: Your input to the tool was invalid')
+			});
+			expect(({} as Record<string, unknown>).polluted).toBeUndefined();
+		});
+
+		test('should bail out on conflicting flattened keys', () => {
+			const conflictTool: vscode.LanguageModelToolInformation = {
+				name: 'conflictTool',
+				description: 'A tool whose flattened input has conflicting paths',
+				inputSchema: {
+					type: 'object',
+					properties: {
+						a: { type: 'object' }
+					},
+					required: ['a']
+				},
+				tags: [],
+				source: undefined
+			};
+
+			(toolsService.tools as vscode.LanguageModelToolInformation[]).push(conflictTool);
+
+			// `a` is both a primitive and a parent of `a.b` — unresolvable.
+			const conflicting = JSON.stringify({
+				'a': 'primitive',
+				'a.b': 'nested'
+			});
+
+			const result = toolsService.validateToolInput('conflictTool', conflicting);
+			expect(result).toMatchObject({
+				error: expect.stringContaining('ERROR: Your input to the tool was invalid')
+			});
+		});
+
+		test('should reject out-of-range array indices in flattened keys', () => {
+			const indexTool: vscode.LanguageModelToolInformation = {
+				name: 'indexTool',
+				description: 'A tool whose flattened input has an enormous array index',
+				inputSchema: {
+					type: 'object',
+					properties: {
+						items: {
+							type: 'array',
+							items: { type: 'string' }
+						}
+					},
+					required: ['items']
+				},
+				tags: [],
+				source: undefined
+			};
+
+			(toolsService.tools as vscode.LanguageModelToolInformation[]).push(indexTool);
+
+			// A huge index would create a massive sparse array; reconstruction
+			// must bail rather than produce one.
+			const huge = JSON.stringify({ 'items[999999999999]': 'value' });
+
+			const result = toolsService.validateToolInput('indexTool', huge);
+			expect(result).toMatchObject({
+				error: expect.stringContaining('ERROR: Your input to the tool was invalid')
+			});
+		});
 	});
 });

extensions/copilot/src/extension/tools/common/toolsService.ts

@@ -130,6 +130,109 @@ function getObjectPropertyByPath(obj: any, jsonPointerPath: string): { parent: a
 	return null;
 }
 
+/**
+ * Property names that must never be used as path segments when reconstructing
+ * objects from untrusted tool input, to avoid prototype pollution.
+ */
+const UNSAFE_PROPERTY_NAMES = new Set(['__proto__', 'constructor', 'prototype']);
+
+/**
+ * Upper bound for array indices accepted when reconstructing flattened tool
+ * input. Caps the reconstructed array length to avoid huge sparse arrays from
+ * untrusted input (e.g. `items[999999999999]`) that would make subsequent Ajv
+ * validation pathologically slow.
+ */
+const MAX_FLATTENED_ARRAY_INDEX = 1000;
+
+/**
+ * Parses a flattened path key (e.g. `questions[0].options[1].label`) into an
+ * ordered list of segments (`['questions', 0, 'options', 1, 'label']`). Object
+ * properties are returned as strings and array indices as numbers. Returns
+ * `undefined` if the key is not a well-formed, contiguous path expression, if
+ * it contains an unsafe property name (e.g. `__proto__`), or if an array index
+ * exceeds {@link MAX_FLATTENED_ARRAY_INDEX}.
+ */
+function parseFlattenedPath(key: string): (string | number)[] | undefined {
+	const segments: (string | number)[] = [];
+	const re = /\.?([^.[\]]+)|\[(\d+)\]/g;
+	let lastIndex = 0;
+	let match: RegExpExecArray | null;
+	while ((match = re.exec(key)) !== null) {
+		// Bail if there is an unexpected character between tokens (e.g. `a..b`).
+		if (match.index !== lastIndex) {
+			return undefined;
+		}
+		if (match[2] !== undefined) {
+			const index = Number(match[2]);
+			// Reject out-of-range indices to avoid huge sparse arrays.
+			if (!Number.isSafeInteger(index) || index > MAX_FLATTENED_ARRAY_INDEX) {
+				return undefined;
+			}
+			segments.push(index);
+		} else {
+			// Reject prototype-pollution keys from untrusted tool input.
+			if (UNSAFE_PROPERTY_NAMES.has(match[1])) {
+				return undefined;
+			}
+			segments.push(match[1]);
+		}
+		lastIndex = re.lastIndex;
+	}
+	if (lastIndex !== key.length || segments.length === 0) {
+		return undefined;
+	}
+	return segments;
+}
+
+/**
+ * Reconstructs a nested object/array structure from an object whose keys are
+ * flattened path expressions. Some models (notably Gemini) serialize nested
+ * tool-call arguments as flat keys like `questions[0].header` instead of a
+ * proper nested object. Returns `undefined` when none of the keys use path
+ * notation (so normal inputs are left untouched), when a key is malformed, or
+ * when keys conflict (e.g. both `a` and `a.b`).
+ */
+function tryUnflattenObject(obj: Record<string, unknown>): Record<string, unknown> | undefined {
+	const keys = Object.keys(obj);
+	if (!keys.some(key => /\.|\[\d+\]/.test(key))) {
+		return undefined;
+	}
+
+	// Use null-prototype containers so untrusted keys cannot reach Object.prototype.
+	const result: Record<string, unknown> = Object.create(null);
+	for (const key of keys) {
+		const path = parseFlattenedPath(key);
+		if (!path) {
+			return undefined;
+		}
+
+		let current: any = result;
+		for (let i = 0; i < path.length - 1; i++) {
+			const segment = path[i];
+			const nextSegment = path[i + 1];
+			const child = current[segment];
+			if (child === undefined) {
+				current[segment] = typeof nextSegment === 'number' ? [] : Object.create(null);
+			} else if (typeof child !== 'object' || child === null) {
+				// Conflicting keys (e.g. both `a` and `a.b`) would require
+				// overwriting a primitive with a container; bail out instead.
+				return undefined;
+			}
+			current = current[segment];
+		}
+
+		const leaf = path[path.length - 1];
+		if (typeof current[leaf] === 'object' && current[leaf] !== null) {
+			// A container already exists at this leaf (e.g. both `a` and `a.b`
+			// where `a` is assigned last); refuse to clobber it.
+			return undefined;
+		}
+		current[leaf] = obj[key];
+	}
+
+	return result;
+}
+
 function ajvValidateForTool(toolName: string, fn: ValidateFunction, inputObj: unknown): IToolValidationResult {
 	// Empty output can be valid when the schema only has optional properties
 	if (fn(inputObj ?? {})) {
@@ -168,6 +271,16 @@ function ajvValidateForTool(toolName: string, fn: ValidateFunction, inputObj: un
 		}
 	}
 
+	// Recovery: some models (notably Gemini) serialize nested arguments as
+	// flattened path keys like `questions[0].header` instead of nested
+	// objects/arrays. Reconstruct the nested structure and re-validate.
+	if (typeof inputObj === 'object' && inputObj !== null && !Array.isArray(inputObj)) {
+		const unflattened = tryUnflattenObject(inputObj as Record<string, unknown>);
+		if (unflattened) {
+			return ajvValidateForTool(toolName, fn, unflattened);
+		}
+	}
+
 	const errors = fn.errors!.map(e => e.message || `${e.instancePath} is invalid}`);
 	return { error: `ERROR: Your input to the tool was invalid (${errors.join(', ')})` };
 }