refactor(util): use structured options for JSON schema cleaning

- Replaced boolean parameters with a `jsonSchemaCleanOptions` struct in `cleanJSONSchema` to improve readability and scalability.
- Updated `CleanJSONSchemaForAntigravity` and related methods to utilize the new options struct.
- Enhanced flexibility for schema transformations with fine-grained control over operations like union flattening, enum type enforcement, and metadata removal.
- Added comprehensive tests to verify correct handling of unions and enum types in schemas.

Closes: #4666

internal/runtime/executor/antigravity_schema_sanitize_test.go

@@ -274,6 +274,46 @@ func TestSanitizeAntigravityRequestSchemasKeepsResponseSchemasPlaceholderFree(t
 	}
 }
 
+func TestSanitizeAntigravityRequestSchemasPreservesResponseUnionAndEnumType(t *testing.T) {
+	payload := `{"request":{
+		"tools":[{"functionDeclarations":[{"name":"tool","parameters":{"type":"object","properties":{
+			"choice":{"anyOf":[{"type":"string"},{"type":"null"}]},
+			"level":{"type":"number","enum":[1,2]}
+		}}}]}],
+		"generationConfig":{"responseSchema":{"type":"object","properties":{
+			"action":{"anyOf":[
+				{"type":"object","properties":{"name":{"type":"string"}},"required":["name"]},
+				{"type":"null"}
+			]},
+			"conviction":{"type":"number","enum":[0.25,0.5,1]}
+		}}}
+	}}`
+
+	got := sanitizeAntigravityRequestSchemas(payload, true)
+	responseSchema := gjson.Get(got, "request.generationConfig.responseSchema")
+	union := responseSchema.Get("properties.action.anyOf")
+	if !union.IsArray() || len(union.Array()) != 2 || union.Get("1.type").String() != "null" {
+		t.Errorf("response anyOf union was flattened: %s", responseSchema.Raw)
+	}
+	conviction := responseSchema.Get("properties.conviction")
+	if gotType := conviction.Get("type").String(); gotType != "number" {
+		t.Errorf("response enum type = %q, want number: %s", gotType, responseSchema.Raw)
+	}
+	for _, enumValue := range conviction.Get("enum").Array() {
+		if enumValue.Type != gjson.String {
+			t.Errorf("response enum value is not a string: %s", conviction.Raw)
+		}
+	}
+
+	toolSchema := gjson.Get(got, "request.tools.0.functionDeclarations.0.parameters")
+	if toolSchema.Get("properties.choice.anyOf").Exists() {
+		t.Errorf("tool anyOf union was not flattened: %s", toolSchema.Raw)
+	}
+	if gotType := toolSchema.Get("properties.level.type").String(); gotType != "string" {
+		t.Errorf("tool enum type = %q, want string: %s", gotType, toolSchema.Raw)
+	}
+}
+
 func TestAntigravityBuildRequestKeepsJSONObjectSchemaPlaceholderFree(t *testing.T) {
 	input := []byte(`{"model":"gemini-3.1-pro-low","messages":[{"role":"user","content":"hi"}],"response_format":{"type":"json_object"}}`)
 	translated := antigravitychat.ConvertOpenAIRequestToAntigravity("gemini-3.1-pro-low", input, false)

internal/util/gemini_schema.go

@@ -24,50 +24,67 @@ const placeholderReasonDescription = "Brief explanation of why you are calling t
 // and replacements such as "enum" and "type" are fabricated. That regression reached production
 // once already; scope every call site to the schema itself.
 
+type jsonSchemaCleanOptions struct {
+	addPlaceholder       bool
+	removeGeminiMetadata bool
+	flattenUnions        bool
+	forceEnumStringType  bool
+}
+
 // CleanJSONSchemaForAntigravity transforms a tool schema to be compatible with Antigravity API.
 // It handles unsupported keywords, type flattening, and schema simplification while preserving
 // semantic information as description hints and adding placeholders required by VALIDATED mode.
 func CleanJSONSchemaForAntigravity(jsonStr string) string {
-	return cleanJSONSchema(jsonStr, true, false)
+	return cleanJSONSchema(jsonStr, jsonSchemaCleanOptions{
+		addPlaceholder:      true,
+		flattenUnions:       true,
+		forceEnumStringType: true,
+	})
 }
 
-// CleanJSONSchemaForAntigravityResponse transforms a response schema without adding tool-only
-// placeholders that would alter the client's structured output contract.
+// CleanJSONSchemaForAntigravityResponse transforms a response schema without applying tool-only
+// compatibility rewrites that would alter the client's structured output contract.
 func CleanJSONSchemaForAntigravityResponse(jsonStr string) string {
-	return cleanJSONSchema(jsonStr, false, false)
+	return cleanJSONSchema(jsonStr, jsonSchemaCleanOptions{})
 }
 
 // CleanJSONSchemaForGemini transforms a JSON schema to be compatible with Gemini tool calling.
 // It removes unsupported keywords and simplifies schemas, without adding empty-schema placeholders.
 func CleanJSONSchemaForGemini(jsonStr string) string {
-	return cleanJSONSchema(jsonStr, false, true)
+	return cleanJSONSchema(jsonStr, jsonSchemaCleanOptions{
+		removeGeminiMetadata: true,
+		flattenUnions:        true,
+		forceEnumStringType:  true,
+	})
 }
 
 // cleanJSONSchema performs the core cleaning operations on the JSON schema.
-func cleanJSONSchema(jsonStr string, addPlaceholder, removeGeminiMetadata bool) string {
+func cleanJSONSchema(jsonStr string, options jsonSchemaCleanOptions) string {
 	// Phase 1: Convert and add hints
 	jsonStr = convertRefsToHints(jsonStr)
 	jsonStr = convertConstToEnum(jsonStr)
-	jsonStr = convertEnumValuesToStrings(jsonStr)
+	jsonStr = convertEnumValuesToStrings(jsonStr, options.forceEnumStringType)
 	jsonStr = addEnumHints(jsonStr)
 	jsonStr = addAdditionalPropertiesHints(jsonStr)
 	jsonStr = moveConstraintsToDescription(jsonStr)
 
 	// Phase 2: Flatten complex structures
 	jsonStr = mergeAllOf(jsonStr)
-	jsonStr = flattenAnyOfOneOf(jsonStr)
+	if options.flattenUnions {
+		jsonStr = flattenAnyOfOneOf(jsonStr)
+	}
 	jsonStr = flattenTypeArrays(jsonStr)
 
 	// Phase 3: Cleanup
 	jsonStr = removeUnsupportedKeywords(jsonStr)
-	if removeGeminiMetadata {
+	if options.removeGeminiMetadata {
 		// Gemini schema cleanup: remove nullable/title and placeholder-only fields.
 		jsonStr = removeKeywords(jsonStr, []string{"nullable", "title"})
 		jsonStr = removePlaceholderFields(jsonStr)
 	}
 	jsonStr = cleanupRequiredFields(jsonStr)
 	// Phase 4: Add placeholder for empty object schemas (Claude VALIDATED mode requirement)
-	if addPlaceholder {
+	if options.addPlaceholder {
 		jsonStr = addEmptySchemaPlaceholder(jsonStr)
 	}
 
@@ -201,9 +218,10 @@ func convertConstToEnum(jsonStr string) string {
 	return jsonStr
 }
 
-// convertEnumValuesToStrings ensures all enum values are strings and the schema type is set to string.
-// Gemini API requires enum values to be of type string, not numbers or booleans.
-func convertEnumValuesToStrings(jsonStr string) string {
+// convertEnumValuesToStrings ensures all enum values use the string representation required by
+// Gemini's proto schema. Tool schemas also require a string type, while response schemas preserve
+// their declared type because the upstream decoder uses it to select the emitted JSON value type.
+func convertEnumValuesToStrings(jsonStr string, forceStringType bool) string {
 	for _, p := range findPaths(jsonStr, "enum") {
 		arr := gjson.Get(jsonStr, p)
 		if !arr.IsArray() {
@@ -215,13 +233,13 @@ func convertEnumValuesToStrings(jsonStr string) string {
 			stringVals = append(stringVals, item.String())
 		}
 
-		// Always update enum values to strings and set type to "string"
-		// This ensures compatibility with Antigravity Gemini which only allows enum for STRING type
 		updated, _ := sjson.SetBytes([]byte(jsonStr), p, stringVals)
 		jsonStr = string(updated)
-		parentPath := trimSuffix(p, ".enum")
-		updated, _ = sjson.SetBytes([]byte(jsonStr), joinPath(parentPath, "type"), "string")
-		jsonStr = string(updated)
+		if forceStringType {
+			parentPath := trimSuffix(p, ".enum")
+			updated, _ = sjson.SetBytes([]byte(jsonStr), joinPath(parentPath, "type"), "string")
+			jsonStr = string(updated)
+		}
 	}
 	return jsonStr
 }

internal/util/gemini_schema_test.go

@@ -764,6 +764,76 @@ func TestCleanJSONSchemaForAntigravityResponseDoesNotAddToolPlaceholders(t *test
 	}
 }
 
+func TestCleanJSONSchemaForAntigravityResponsePreservesUnions(t *testing.T) {
+	input := `{
+		"type":"object",
+		"properties":{
+			"action":{"anyOf":[
+				{"type":"object","properties":{"name":{"type":"string"}},"required":["name"]},
+				{"type":"null"}
+			]},
+			"label":{"oneOf":[{"type":"string"},{"type":"null"}]}
+		}
+	}`
+
+	result := gjson.Parse(CleanJSONSchemaForAntigravityResponse(input))
+	for _, testCase := range []struct {
+		path      string
+		wantTypes []string
+	}{
+		{path: "properties.action.anyOf", wantTypes: []string{"object", "null"}},
+		{path: "properties.label.oneOf", wantTypes: []string{"string", "null"}},
+	} {
+		union := result.Get(testCase.path)
+		if !union.IsArray() {
+			t.Errorf("response union %s was flattened: %s", testCase.path, result.Raw)
+			continue
+		}
+		var gotTypes []string
+		for _, branch := range union.Array() {
+			gotTypes = append(gotTypes, branch.Get("type").String())
+		}
+		if !reflect.DeepEqual(gotTypes, testCase.wantTypes) {
+			t.Errorf("response union %s types = %v, want %v: %s", testCase.path, gotTypes, testCase.wantTypes, result.Raw)
+		}
+	}
+}
+
+func TestCleanJSONSchemaForAntigravityResponsePreservesEnumType(t *testing.T) {
+	input := `{
+		"type":"object",
+		"properties":{
+			"conviction":{"type":"number","enum":[0.25,0.5,1]},
+			"count":{"type":"integer","enum":[1,2]}
+		}
+	}`
+
+	result := gjson.Parse(CleanJSONSchemaForAntigravityResponse(input))
+	for _, testCase := range []struct {
+		path       string
+		wantType   string
+		wantValues []string
+	}{
+		{path: "properties.conviction", wantType: "number", wantValues: []string{"0.25", "0.5", "1"}},
+		{path: "properties.count", wantType: "integer", wantValues: []string{"1", "2"}},
+	} {
+		schema := result.Get(testCase.path)
+		if gotType := schema.Get("type").String(); gotType != testCase.wantType {
+			t.Errorf("%s type = %q, want %q: %s", testCase.path, gotType, testCase.wantType, result.Raw)
+		}
+		var gotValues []string
+		for _, enumValue := range schema.Get("enum").Array() {
+			if enumValue.Type != gjson.String {
+				t.Errorf("%s enum value is not a string: %s", testCase.path, enumValue.Raw)
+			}
+			gotValues = append(gotValues, enumValue.String())
+		}
+		if !reflect.DeepEqual(gotValues, testCase.wantValues) {
+			t.Errorf("%s enum values = %v, want %v: %s", testCase.path, gotValues, testCase.wantValues, result.Raw)
+		}
+	}
+}
+
 // ============================================================================
 // Format field handling (ad-hoc patch removal)
 // ============================================================================
@@ -862,6 +932,14 @@ func TestCleanJSONSchemaForAntigravity_NumericEnumToString(t *testing.T) {
 	}`
 
 	result := CleanJSONSchemaForAntigravity(input)
+	parsed := gjson.Parse(result)
+
+	// Tool enum schemas require both string values and a string type.
+	for _, path := range []string{"properties.priority", "properties.level"} {
+		if gotType := parsed.Get(path + ".type").String(); gotType != "string" {
+			t.Errorf("Tool enum type at %s = %q, want string: %s", path, gotType, result)
+		}
+	}
 
 	// Numeric enum values should be converted to strings
 	if strings.Contains(result, `"enum":[0,1,2]`) {