diff --git a/core/plugins/toolprune/plugin.go b/core/plugins/toolprune/plugin.go index 95ec37a52..73a9f4ab8 100644 --- a/core/plugins/toolprune/plugin.go +++ b/core/plugins/toolprune/plugin.go @@ -390,22 +390,39 @@ func (p *ToolPrune) OnRequest(_ context.Context, pctx *pipeline.Context) (action break } } - if orphanedCacheControl.Exists() && lastSurvivor >= 0 && - !gjson.GetBytes(body, fmt.Sprintf("tools.%d.cache_control", lastSurvivor)).Exists() { - if out, err = sjson.SetRawBytes(out, - fmt.Sprintf("tools.%d.cache_control", lastSurvivor), - []byte(orphanedCacheControl.Raw)); err != nil { - slog.Warn("tool-prune: could not preserve cache_control, forwarding original", "err", err) - return action + // One splice, not one per victim: sjson.DeleteBytes returns a fresh copy + // of the whole body on every call, so k victims used to cost k+1 full + // copies — measured at 41% of a live proxy's total allocations. Each + // survivor's bytes are kept verbatim; only the separators between them + // normalize to ",", which no cache prefix can have relied on anyway — + // pruning already rewrote the array. + joined := make([]byte, 0, len(body)) + joined = append(joined, '[') + for i := range raw { + if victimSet[i] { + continue } - slog.Debug("tool-prune: moved cache_control to the last surviving tool", "index", lastSurvivor) - } - // Descending, so an earlier deletion never shifts a later index. - for i := len(victims) - 1; i >= 0; i-- { - if out, err = sjson.DeleteBytes(out, fmt.Sprintf("tools.%d", victims[i])); err != nil { - slog.Warn("tool-prune: delete failed, forwarding original", "index", victims[i], "err", err) - return action + if len(joined) > 1 { + joined = append(joined, ',') } + survivor := raw[i].Raw + if i == lastSurvivor && orphanedCacheControl.Exists() && + !gjson.GetBytes(body, fmt.Sprintf("tools.%d.cache_control", i)).Exists() { + patched, err := sjson.SetRawBytes([]byte(survivor), "cache_control", []byte(orphanedCacheControl.Raw)) + if err != nil { + slog.Warn("tool-prune: could not preserve cache_control, forwarding original", "err", err) + return action + } + survivor = string(patched) + slog.Debug("tool-prune: moved cache_control to the last surviving tool", "index", i) + } + joined = append(joined, survivor...) + } + joined = append(joined, ']') + var err error + if out, err = sjson.SetRawBytes(out, "tools", joined); err != nil { + slog.Warn("tool-prune: rewrite of the tools array failed, forwarding original", "err", err) + return action } } if len(out) >= len(body) { diff --git a/core/plugins/toolprune/plugin_test.go b/core/plugins/toolprune/plugin_test.go index 89c28a96a..f7c7bbd8a 100644 --- a/core/plugins/toolprune/plugin_test.go +++ b/core/plugins/toolprune/plugin_test.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "fmt" + "runtime" "strings" "sync" "testing" @@ -97,10 +98,11 @@ func TestPrune_LeavesEveryOtherByteIntact(t *testing.T) { } } -// TestPrune_DescendingDeletion: removing several tools by index only works if -// the deletions run high-to-low. An ascending loop would shift the array under -// itself and delete the wrong elements — here it would leave "Bash" and remove -// something else, so the assertion catches exactly that bug. +// TestPrune_DescendingDeletion: removing several tools in one pass must remove +// exactly the named ones — this test dates from when the rewrite deleted +// elements by index, where an ascending loop shifted the array under itself +// and removed the wrong tools. The name keeps the history; the assertion still +// pins multi-victim selection. func TestPrune_DescendingDeletion(t *testing.T) { body := `{"tools":[{"name":"A"},{"name":"B"},{"name":"C"},{"name":"D"},{"name":"E"}]}` p := configured(t, "A", "B", "D") @@ -1252,3 +1254,38 @@ func TestPrune_SettlementCalibratesOnTheBytesALaterWriterSent(t *testing.T) { pricing.EstimateTokensFromBytes(ev.BytesRemoved, prompt, ev.BodyBytesAfter)) } } + +// TestPrune_RewriteDoesNotCopyTheBodyPerVictim: the rewrite used to splice the +// tools array with one sjson.DeleteBytes per victim, and each of those returns +// a fresh copy of the whole body — k pruned tools cost k+1 full copies, +// measured at 41% of a live proxy's total allocations. The splice must not +// scale with the number of victims. +func TestPrune_RewriteDoesNotCopyTheBodyPerVictim(t *testing.T) { + pad := strings.Repeat("x", 256<<10) + var b strings.Builder + b.WriteString(`{"tools":[`) + for i, name := range []string{"A", "B", "C", "D", "E", "F", "G", "H", "I", "J"} { + if i > 0 { + b.WriteByte(',') + } + fmt.Fprintf(&b, `{"name":%q,"pad":"%s"}`, name, pad) + } + b.WriteString(`],"messages":[]}`) + body := b.String() + + p := configured(t, "A", "B", "C", "D", "E", "F", "G", "H", "I") // 9 of 10 + pctx := inferenceCtx("/v1/messages", body, "A", "B", "C", "D", "E", "F", "G", "H", "I", "J") + + runtime.GC() + var before, after runtime.MemStats + runtime.ReadMemStats(&before) + run(t, p, pctx) + runtime.ReadMemStats(&after) + + if !pctx.BodyMutated() { + t.Fatal("expected a rewrite") + } + if got, want := after.TotalAlloc-before.TotalAlloc, uint64(3*len(body)); got > want { + t.Errorf("pruning 9 of 10 tools allocated %d bytes for a %d-byte body (want <= %d): the rewrite copies the body once per victim", got, len(body), want) + } +}