Skip to content

Perf: Splice tool-prune's tools array in one pass #1606

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
huang195 merged 1 commit into from
Oct 30, 2026
Merged

Perf: Splice tool-prune's tools array in one pass #1606

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
45 changes: 31 additions & 14 deletions core/plugins/toolprune/plugin.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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))

@clawgenti clawgenti Oct 30, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

suggestion: The PR description defers this cap as "over-reserves when most of a large manifest is pruned, low value". The bigger cost is elsewhere: len(body) includes messages, and on long conversations the history is usually much larger than the tool manifest. In that case this scratch buffer is a second body-sized allocation, which cancels most of the saving when only a few tools are pruned.

I measured it with a 4 MB history, 10 x 10 KB tools, and 2 victims:

TotalAlloc / body
old (DeleteBytes loop) 3.07x
this PR (cap = len(body)) 3.04x
cap = len(tools.Raw) 2.07x 
joined := make([]byte, 0, len(tools.Raw))

tools.Raw is always at least the size of the survivors plus brackets and commas. The only exception is the cache_control patch adding a few bytes, and append handles that. The full suite passes with this change. It's a one-line fix, so it seems worth doing here rather than deferring.

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() {

@clawgenti clawgenti Oct 30, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: raw[i].Get("cache_control").Exists() gives the same answer from the element you already hold. gjson.GetBytes(body, "tools.%d.cache_control") re-walks the body up to index i, plus a Sprintf. It only runs once per request, so this is cosmetic, but it fits the spirit of the PR.

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) {
Expand Down
45 changes: 41 additions & 4 deletions core/plugins/toolprune/plugin_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import (
"context"
"encoding/json"
"fmt"
"runtime"
"strings"
"sync"
"testing"
Expand Down Expand Up @@ -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")
Expand Down Expand Up @@ -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)

@clawgenti clawgenti Oct 30, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: Every byte of this body is in tools, so len(body) and len(tools.Raw) are almost equal and the test can't tell the two buffer sizings apart (see the comment on plugin.go:399). If you add a large messages pad (for example 2-4 MB of history) and keep the victim count low, the test models the common Claude Code case. With the tools.Raw cap you could then tighten the bound to about 2*len(body) plus slack.

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