From 2b63d6bcda136af1d3638be8e0038658911fb217 Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Thu, 30 Jul 2026 03:39:48 +0800 Subject: [PATCH] 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 --- .../antigravity_schema_sanitize_test.go | 40 ++++++++++ internal/util/gemini_schema.go | 54 ++++++++----- internal/util/gemini_schema_test.go | 78 +++++++++++++++++++ 3 files changed, 154 insertions(+), 18 deletions(-) diff --git a/internal/runtime/executor/antigravity_schema_sanitize_test.go b/internal/runtime/executor/antigravity_schema_sanitize_test.go index 6151ae4c4..139374bb7 100644 --- a/internal/runtime/executor/antigravity_schema_sanitize_test.go +++ b/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) diff --git a/internal/util/gemini_schema.go b/internal/util/gemini_schema.go index 51a414f79..f3b99d295 100644 --- a/internal/util/gemini_schema.go +++ b/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 } diff --git a/internal/util/gemini_schema_test.go b/internal/util/gemini_schema_test.go index 20d10b4d1..50a66bbdd 100644 --- a/internal/util/gemini_schema_test.go +++ b/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]`) {