feat(runtime): optimize antigravity schema sanitization to scale linearly with declarations

- Refactor antigravity tool schema sanitization to process declarations and elements in a single pass.
- Avoid quadratic whole-document rewrites for payloads with many function declarations.
- Preserve original array whitespace and untouched elements during byte splicing.

Closes: #6714

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

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)
+	}
+}