diff --git a/internal/runtime/executor/antigravity_executor_request.go b/internal/runtime/executor/antigravity_executor_request.go index 3c25099efd..51c396f741 100644 --- a/internal/runtime/executor/antigravity_executor_request.go +++ b/internal/runtime/executor/antigravity_executor_request.go @@ -213,20 +213,20 @@ func sanitizeAntigravityToolSchemaDocument(payloadStr string, useAntigravitySche return sanitizeAntigravityToolSchemaDocumentForModel(payloadStr, useAntigravitySchema, false) } -func sanitizeAntigravityToolSchemaDocumentForModel(payloadStr string, useAntigravitySchema, isClaudeModel bool) string { - for _, base := range antigravityFunctionDeclarationPaths(payloadStr) { - oldPath := base + ".parametersJsonSchema" - if !gjson.Get(payloadStr, oldPath).Exists() { - continue - } - renamed, errRename := util.RenameKey(payloadStr, oldPath, base+".parameters") - if errRename != nil { - log.Debugf("antigravity: failed to rename %s: %v", oldPath, errRename) - continue - } - payloadStr = renamed - } +// antigravityDeclarationContainers lists both spellings of the declarations +// array a tool may carry. +var antigravityDeclarationContainers = []string{"functionDeclarations", "function_declarations"} +// sanitizeAntigravityToolSchemaDocumentForModel renames and cleans every +// declaration schema in the tools-only document. Each declaration is edited +// on its own raw bytes and the results are spliced back into the tools array +// in a single pass, so the cost stays linear in the document size instead of +// scaling with declarations x document size (Issue 6714). The emitted bytes +// match the previous whole-document edits: changed elements are spliced at +// their original offsets, so untouched bytes (including whitespace around +// array elements) survive verbatim, and changed nodes go through the same +// sjson edits, scoped to one declaration. +func sanitizeAntigravityToolSchemaDocumentForModel(payloadStr string, useAntigravitySchema, isClaudeModel bool) string { toolSchemaCleaner := func(schema string, isParameterSchema bool) string { // Only declaration parameter schemas become Claude custom.input_schema upstream, // so only they need Claude union flattening. Function return schemas keep the @@ -241,23 +241,163 @@ func sanitizeAntigravityToolSchemaDocumentForModel(payloadStr string, useAntigra return toolSchemaCleaner(schema, isParameterSchema) }, schemaRaw) } - for _, schemaPath := range antigravityDeclarationSchemaPaths(payloadStr) { - schema := gjson.Get(payloadStr, schemaPath) - if !schema.Exists() { + + tools := gjson.Parse(payloadStr).Get("request.tools") + if !tools.IsArray() { + return payloadStr + } + elements := make([]gjson.Result, 0) + edits := make([]antigravityElementEdit, 0) + tools.ForEach(func(_, tool gjson.Result) bool { + elements = append(elements, tool) + if !tool.IsObject() { + return true + } + toolRaw := tool.Raw + for _, declKey := range antigravityDeclarationContainers { + decls := tool.Get(declKey) + if !decls.IsArray() { + continue + } + cleanedDecls, declsChanged := sanitizeAntigravityDeclarations(decls, cleanNestedToolSchema) + if !declsChanged { + continue + } + updated, errSet := sjson.SetRawBytes([]byte(toolRaw), declKey, []byte(cleanedDecls)) + if errSet != nil { + log.Debugf("antigravity: failed to write cleaned %s: %v", declKey, errSet) + continue + } + toolRaw = string(updated) + } + if toolRaw != tool.Raw { + edits = append(edits, antigravityElementEdit{offset: tool.Index - tools.Index, raw: tool.Raw, replacement: toolRaw}) + } + return true + }) + if len(edits) == 0 { + return payloadStr + } + updated, errSet := sjson.SetRawBytes([]byte(payloadStr), "request.tools", []byte(applyAntigravityElementEdits(tools.Raw, tools.Index, elements, edits))) + if errSet != nil { + log.Debugf("antigravity: failed to write cleaned request.tools: %v", errSet) + return payloadStr + } + return string(updated) +} + +// antigravityElementEdit replaces one array element (raw, at offset within +// the array's raw bytes) with replacement. +type antigravityElementEdit struct { + offset int + raw string + replacement string +} + +// applyAntigravityElementEdits splices element replacements into an array's +// raw bytes in offset order, keeping every other byte (commas, whitespace, +// unchanged elements) verbatim. Edits whose offsets do not line up with the +// array bytes fall back to joining the elements with commas: still valid +// and semantically identical, just without the original separators. gjson +// reports element offsets into the parsed document, so offsets here are +// relative to the array start. +func applyAntigravityElementEdits(arrayRaw string, arrayIndex int, elements []gjson.Result, edits []antigravityElementEdit) string { + var rebuilt strings.Builder + rebuilt.Grow(len(arrayRaw)) + cursor := 0 + for _, edit := range edits { + if edit.offset < cursor || edit.offset+len(edit.raw) > len(arrayRaw) || arrayRaw[edit.offset:edit.offset+len(edit.raw)] != edit.raw { + return joinAntigravityElements(arrayIndex, elements, edits) + } + rebuilt.WriteString(arrayRaw[cursor:edit.offset]) + rebuilt.WriteString(edit.replacement) + cursor = edit.offset + len(edit.raw) + } + rebuilt.WriteString(arrayRaw[cursor:]) + return rebuilt.String() +} + +// joinAntigravityElements rebuilds an array by joining element bytes with +// commas, substituting replacements where recorded. It is the fallback for +// applyAntigravityElementEdits when splice offsets cannot be trusted. +func joinAntigravityElements(arrayIndex int, elements []gjson.Result, edits []antigravityElementEdit) string { + replacements := make(map[int]string, len(edits)) + for _, edit := range edits { + replacements[edit.offset] = edit.replacement + } + var rebuilt strings.Builder + rebuilt.WriteByte('[') + for i, element := range elements { + if i > 0 { + rebuilt.WriteByte(',') + } + if replacement, ok := replacements[element.Index-arrayIndex]; ok { + rebuilt.WriteString(replacement) + } else { + rebuilt.WriteString(element.Raw) + } + } + rebuilt.WriteByte(']') + return rebuilt.String() +} + +// sanitizeAntigravityDeclarations rebuilds one declarations array, applying +// the rename and schema cleaning to each declaration in isolation. It +// reports whether any declaration actually changed so callers can keep the +// original bytes when the array is already clean. +func sanitizeAntigravityDeclarations(decls gjson.Result, clean func(schemaRaw string, isParameterSchema bool) string) (string, bool) { + elements := make([]gjson.Result, 0) + edits := make([]antigravityElementEdit, 0) + decls.ForEach(func(_, decl gjson.Result) bool { + elements = append(elements, decl) + if !decl.IsObject() { + return true + } + cleanedDecl := sanitizeAntigravityDeclaration(decl, clean) + if cleanedDecl != decl.Raw { + edits = append(edits, antigravityElementEdit{offset: decl.Index - decls.Index, raw: decl.Raw, replacement: cleanedDecl}) + } + return true + }) + if len(edits) == 0 { + return "", false + } + return applyAntigravityElementEdits(decls.Raw, decls.Index, elements, edits), true +} + +// sanitizeAntigravityDeclaration applies the declaration-level edits that +// used to run against the whole tools document: parametersJsonSchema is +// renamed onto parameters first, then every object schema is cleaned in +// place. Failed edits keep the declaration as-is and are only logged. +func sanitizeAntigravityDeclaration(decl gjson.Result, clean func(schemaRaw string, isParameterSchema bool) string) string { + declRaw := decl.Raw + if decl.Get("parametersJsonSchema").Exists() { + renamed, errRename := util.RenameKey(declRaw, "parametersJsonSchema", "parameters") + if errRename != nil { + log.Debugf("antigravity: failed to rename parametersJsonSchema: %v", errRename) + } else { + declRaw = renamed + decl = gjson.Parse(declRaw) + } + } + for _, key := range antigravityDeclarationSchemaKeys { + schema := decl.Get(key) + if !schema.IsObject() { continue } - cleanedSchema := cleanNestedToolSchema(schema.Raw, isParameterSchemaPath(schemaPath)) + cleanedSchema := clean(schema.Raw, isParameterSchemaPath(key)) if cleanedSchema == schema.Raw { continue } - updated, errSet := sjson.SetRawBytes([]byte(payloadStr), schemaPath, []byte(cleanedSchema)) + updated, errSet := sjson.SetRawBytes([]byte(declRaw), key, []byte(cleanedSchema)) if errSet != nil { - log.Debugf("antigravity: failed to write cleaned schema at %s: %v", schemaPath, errSet) + log.Debugf("antigravity: failed to write cleaned schema at %s: %v", key, errSet) continue } - payloadStr = string(updated) + declRaw = string(updated) + decl = gjson.Parse(declRaw) } - return payloadStr + return declRaw } // sanitizeAntigravityGenerationSchemas batches every schema edit within one diff --git a/internal/runtime/executor/antigravity_schema_sanitize_test.go b/internal/runtime/executor/antigravity_schema_sanitize_test.go index abfbf4befc..e7736478a8 100644 --- a/internal/runtime/executor/antigravity_schema_sanitize_test.go +++ b/internal/runtime/executor/antigravity_schema_sanitize_test.go @@ -2,6 +2,8 @@ package executor import ( "encoding/json" + "fmt" + "runtime" "strings" "testing" @@ -904,3 +906,95 @@ func TestAntigravityBuildRequestFlattensToolUnionsForClaude_Issue3662Regression( } }) } + +// TestSanitizeAntigravityRequestSchemasScalesLinearlyWithDeclarations is the +// regression test for the quadratic tool sanitizing reported in Issue 6714: +// sanitizing must not rewrite the whole tools document once per function +// declaration. The allocation bound is asserted directly in the test (not via +// a benchmark) so a return to per-declaration document rewrites fails loudly. +func TestSanitizeAntigravityRequestSchemasScalesLinearlyWithDeclarations(t *testing.T) { + const declCount = 256 + + var input, want strings.Builder + input.WriteString(`{"request":{"tools":[{"functionDeclarations":[`) + want.WriteString(`{"request":{"tools":[{"functionDeclarations":[`) + for i := 0; i < declCount; i++ { + if i > 0 { + input.WriteByte(',') + want.WriteByte(',') + } + // parametersJsonSchema forces the rename onto parameters and $id forces + // a cleaner rewrite of the schema. The padded description sits on the + // declaration, outside the schema, so whole-document copies must pay for + // it while per-declaration edits do not; the bound below then measures + // document-rewrite cost only, not the cleaner's per-schema cost (Issue 6713). + fmt.Fprintf(&input, `{"name":"tool_%d","description":"%s","parametersJsonSchema":{"type":"object","$id":"drop","properties":{"q":{"type":"string"}}}}`, i, strings.Repeat("x", 4096)) + // The schema has no required properties, so the Claude VALIDATED + // placeholder is inserted even on the Gemini path (antigravity=true). + fmt.Fprintf(&want, `{"name":"tool_%d","description":"%s","parameters":{"type":"object","properties":{"q":{"type":"string"},"_":{"type":"boolean"}},"required":["_"]}}`, i, strings.Repeat("x", 4096)) + } + input.WriteString(`]}]}}`) + want.WriteString(`]}]}}`) + + payload := input.String() + var before, after runtime.MemStats + runtime.GC() + runtime.ReadMemStats(&before) + got := sanitizeAntigravityRequestSchemasForModel(payload, true, false) + runtime.ReadMemStats(&after) + + if got != want.String() { + t.Fatalf("sanitized tools diverged from the expected bytes\nfirst declaration want: %s\nfirst declaration got: %s", + gjson.Get(want.String(), "request.tools.0.functionDeclarations.0").Raw, + gjson.Get(got, "request.tools.0.functionDeclarations.0").Raw) + } + + // Whole-document work is bounded by a constant number of passes: extracting + // request.tools, rebuilding it, and writing it back. Everything beyond a + // small multiple of the payload size means per-declaration document rewrites + // are back; the pre-fix implementation allocated O(declCount x payload). + const allowedCopies = 32 + if allocated := after.TotalAlloc - before.TotalAlloc; allocated > uint64(allowedCopies*len(payload)) { + t.Fatalf("sanitizing %d declarations allocated %d bytes for a %d byte payload, want at most %d (cost must be linear in the tools size, not per declaration)", + declCount, allocated, len(payload), allowedCopies*len(payload)) + } +} + +// TestSanitizeAntigravityRequestSchemasPreservesToolArrayWhitespace pins the +// byte-equivalence claim of the per-declaration rewrite: rebuilding the tools +// and declarations arrays must keep the whitespace around array elements, so +// a pretty-printed payload only changes inside the cleaned declarations. The +// expected value is the input with only the affected schema bytes replaced, +// exactly what the previous whole-document sjson edits produced. +func TestSanitizeAntigravityRequestSchemasPreservesToolArrayWhitespace(t *testing.T) { + input := "{\n" + + " \"request\": {\n" + + " \"tools\": [\n" + + " {\"functionDeclarations\": [\n" + + " {\"name\": \"a\", \"parametersJsonSchema\": {\"type\": \"object\", \"$id\": \"drop\", \"properties\": {\"q\": {\"type\": \"string\"}}}} ,\n" + + " {\"name\": \"b\", \"parameters\": {\"type\": \"object\", \"required\": [\"q\"], \"properties\": {\"q\": {\"type\": \"string\"}}}} ,\n" + + " {\"name\": \"c\"}\n" + + " ]} ,\n" + + " {\"googleSearch\": {}}\n" + + " ]\n" + + " }\n" + + "}\n" + + // Declaration a is renamed and cleaned in place (sjson compacts the + // replacement gaps inside the edited node, as the whole-document edits + // did). Declaration b is already clean, and declarations b and c and + // the googleSearch tool are untouched, so every other byte (including the + // whitespace around commas and brackets) must survive verbatim. The odd + // "}} ," separator is the declaration's own closing bytes, one space of + // inter-element whitespace, the comma, and the newline. + want := strings.Replace(input, + `{"name": "a", "parametersJsonSchema": {"type": "object", "$id": "drop", "properties": {"q": {"type": "string"}}}}`, + `{"name": "a","parameters":{"type": "object", "properties": {"q": {"type": "string"},"_":{"type":"boolean"}},"required":["_"]}}`, 1) + if want == input { + t.Fatal("test fixture is broken: declaration a was not replaced in the expected value") + } + + if got := sanitizeAntigravityRequestSchemasForModel(input, true, false); got != want { + t.Fatalf("pretty-printed tools lost whitespace outside the cleaned declaration\nwant:\n%s\ngot:\n%s", want, got) + } +}