diff --git a/components/offload/prefixpin.go b/components/offload/prefixpin.go new file mode 100644 index 0000000..9861372 --- /dev/null +++ b/components/offload/prefixpin.go @@ -0,0 +1,232 @@ +package offload + +import ( + "encoding/json" + "strconv" + "strings" + + bschemas "github.com/maximhq/bifrost/core/schemas" + "github.com/rossoctl/context-guru/components" + "github.com/rossoctl/context-guru/schema" + "gopkg.in/yaml.v3" +) + +func init() { components.Register("prefixpin", newPrefixpin) } + +// Prefixpin restores prefix stability when an agent REWRITES an early message on +// every turn. It is the counterpart to cacheinject: cacheinject optimises where +// the cache boundary goes, prefixpin fixes the case where no boundary can help. +// +// Why this exists, from measurement rather than theory. Across 1,955 real Bob +// requests on SWE-bench, 98.0% of turns were append-only and cached fine. The +// other 2% cost 5,796,220 tokens — 71.8% of ALL uncached input — because a single +// early message mutated each turn. On one task the agent re-emitted a running +// / at message index 1, re-rendering an iteration +// counter in ~20 places ("THIRTY-SECOND" -> "THIRTY-THIRD", "thirty-second" -> +// "thirty-third", "32" -> "33"): 152 changed characters out of 6,024, i.e. 98.5% +// identical content, sitting ~1,374 tokens into a 181k-token prefix. Only 0.76% of +// the prefix survived and the cache hit rate collapsed from 98% to 5.7%. +// +// The economics are lopsided. A prompt-cache read costs 0.1x base input and +// uncached input costs 1.0x, so a mutation below the cache boundary makes every +// token above it TEN TIMES more expensive. Compare placement tuning, which only +// ever moves tokens between read (0.1x) and cache-write (1.25x). That is why this +// is worth ~31% of Bob's input cost while every placement change measured ~0%. +// +// How: for each early message, remember the first text seen for its (index, role) +// slot in this session. On a later turn, if that slot's text has CHANGED but is +// still recognisably the same content (same shape, high similarity), re-send the +// FIRST rendering so the prefix stays byte-identical. +// +// LOSSINESS — this is an Offload, deliberately, not a Reformat. The model sees the +// pinned (older) text rather than what the agent just wrote, so information IS +// withheld: a counter reads stale. That is a real behavioural change and the +// reason for the guards below. The original is stashed so `expand` can recover it. +// +// Guards, each closing a way this could do harm: +// - only messages at index < MaxPinIndex (an early, structural slot; never the +// working tail the agent is actively reasoning about) +// - only when similarity >= MinSimilarity (a genuinely rewritten-in-place block, +// not a different message that happens to occupy the slot) +// - only after the slot has churned RepeatThreshold times, so a one-off edit is +// never pinned — only a per-turn churn pattern +// - never on the newest message, never on tool results +type Prefixpin struct { + // MaxPinIndex bounds pinning to structurally-early messages. 0 disables. + MaxPinIndex int `yaml:"max_pin_index"` + // MinSimilarity is the character-level overlap required to treat a changed slot + // as the same content rewritten in place. + MinSimilarity float64 `yaml:"min_similarity"` + // RepeatThreshold is how many times a slot must churn before pinning starts. + RepeatThreshold int `yaml:"repeat_threshold"` + // MinTokens skips slots too small to be worth the behavioural risk. + MinTokens int `yaml:"min_tokens"` +} + +func newPrefixpin(raw []byte) (components.Component, error) { + p := &Prefixpin{MaxPinIndex: 4, MinSimilarity: 0.80, RepeatThreshold: 2, MinTokens: 200} + if len(raw) > 0 { + if err := yaml.Unmarshal(raw, p); err != nil { + return nil, err + } + } + return p, nil +} + +func (Prefixpin) Name() string { return "prefixpin" } + +// Enabled everywhere: the failure mode is provider-independent. It bites hardest on +// implicit-cache backends (Gemini/Bob, OpenAI) where there is no cache_control to +// place and prefix stability is the ONLY available lever. +func (p *Prefixpin) Enabled(c *components.Ctx) bool { return p.MaxPinIndex > 0 } + +func (p *Prefixpin) Offload(req *bschemas.BifrostChatRequest, rep *components.Report, c *components.Ctx) ([]string, error) { + if c == nil || c.Store == nil || c.Session == "" || len(req.Input) < 2 { + rep.Skipped = true + return nil, nil + } + + limit := p.MaxPinIndex + if limit > len(req.Input)-1 { + limit = len(req.Input) - 1 // never the newest message + } + + var keys []string + acted := false + for i := 0; i < limit; i++ { + m := &req.Input[i] + if !schema.Rewritable(*m) || m.Role == bschemas.ChatMessageRoleTool { + continue + } + cur := schema.MessageText(*m) + if cur == "" || schema.TextTokens(cur) < p.MinTokens { + continue + } + if skipReduce(c, cur) { + continue // already offloaded, or the agent expanded it + } + + st := p.loadSlot(c, i, string(m.Role)) + if st == nil { + p.saveSlot(c, i, string(m.Role), &slotState{First: cur, Churn: 0}) + continue + } + if st.First == cur { + continue // stable already: nothing to do, and nothing to pay for + } + if similarity(st.First, cur) < p.MinSimilarity { + // A different message now occupies this slot (the transcript was + // restructured, not rewritten). Re-baseline rather than pin the wrong text. + p.saveSlot(c, i, string(m.Role), &slotState{First: cur, Churn: 0}) + continue + } + + st.Churn++ + p.saveSlot(c, i, string(m.Role), st) + if st.Churn < p.RepeatThreshold { + continue // one-off edit; not yet evidence of per-turn churn + } + + // Pin: re-send the first rendering so the prefix hashes identically. The + // current text is stashed under the marker key so expand can recover it. + k := pinKey(c.Session, i) + c.Store.Put(k, []byte(cur)) + schema.SetMessageText(m, st.First) + keys = append(keys, k) + acted = true + } + + if !acted { + rep.Skipped = true + } + return keys, nil +} + +// --------------------------------------------------------------------------- // + +type slotState struct { + First string `json:"first"` + Churn int `json:"churn"` +} + +func pinKey(session string, i int) string { + return "cg:pin:" + session + ":" + strconv.Itoa(i) +} + +func slotKey(session string, i int, role string) string { + return "cg:pinslot:" + session + ":" + strconv.Itoa(i) + ":" + role +} + +func (p *Prefixpin) loadSlot(c *components.Ctx, i int, role string) *slotState { + b, ok := c.Store.Get(slotKey(c.Session, i, role)) + if !ok || len(b) == 0 { + return nil + } + var s slotState + if json.Unmarshal(b, &s) != nil { + return nil + } + return &s +} + +func (p *Prefixpin) saveSlot(c *components.Ctx, i int, role string, s *slotState) { + if b, err := json.Marshal(s); err == nil { + c.Store.Put(slotKey(c.Session, i, role), b) + } +} + +// similarity estimates content overlap in [0,1] via line-shingle containment. +// +// A prefix+suffix overlap ratio is NOT adequate here, and the traces show exactly +// why: the real churning block differed by 152 characters out of 6,024 (98.5% +// identical by edit distance) but the edits were scattered over 20 separate hunks +// — a repeated counter rendered in several places ("THIRTY-SECOND"/"thirty-second" +// /"32"). Any one early hunk truncates the common prefix and any one late hunk +// truncates the common suffix, so that measure scored it 0.075 and the guard +// rejected the very case it exists to catch. +// +// Line shingles are immune to scattered edits (only the lines containing an edit +// are lost), are O(n) with a bounded map, and need no quadratic edit distance on a +// message that may be 100k tokens. +func similarity(a, b string) float64 { + if a == b { + return 1 + } + if a == "" || b == "" { + return 0 + } + al, bl := strings.Split(a, "\n"), strings.Split(b, "\n") + // Guard the map size on pathological input (one enormous line, or a million + // lines): both extremes fall back to the cheap character ratio, which is only + // used to reject wildly different content. + if len(al) < 4 || len(bl) < 4 { + la, lb := float64(len(a)), float64(len(b)) + if lb > la { + la, lb = lb, la + } + return lb / la + } + set := make(map[string]int, len(al)) + for _, l := range al { + l = strings.TrimSpace(l) + if l != "" { + set[l]++ + } + } + hit, total := 0, 0 + for _, l := range bl { + l = strings.TrimSpace(l) + if l == "" { + continue + } + total++ + if set[l] > 0 { + set[l]-- + hit++ + } + } + if total == 0 { + return 0 + } + return float64(hit) / float64(total) +} diff --git a/components/offload/prefixpin_test.go b/components/offload/prefixpin_test.go new file mode 100644 index 0000000..26090db --- /dev/null +++ b/components/offload/prefixpin_test.go @@ -0,0 +1,263 @@ +package offload + +import ( + "strconv" + "strings" + "testing" + + bschemas "github.com/maximhq/bifrost/core/schemas" + "github.com/rossoctl/context-guru/components" + "github.com/rossoctl/context-guru/schema" + "github.com/rossoctl/context-guru/store" +) + +func pinCtx() *components.Ctx { + return &components.Ctx{Session: "s1", Store: store.NewMemory(store.Options{})} +} + +func msg(role bschemas.ChatMessageRole, text string) bschemas.ChatMessage { + m := bschemas.ChatMessage{Role: role} + schema.SetMessageText(&m, text) + return m +} + +// A big block whose only per-turn change is a counter — the exact shape observed in +// the Bob traces ("the THIRTY-SECOND time" -> "THIRTY-THIRD"). +func scratchpad(n string) string { + return "\nThe user has given the same instruction for the " + n + + " time.\n" + strings.Repeat("stable analysis line that does not change.\n", 200) +} + +func turn(t *testing.T, p *Prefixpin, c *components.Ctx, msgs []bschemas.ChatMessage) ([]bschemas.ChatMessage, *components.Report) { + t.Helper() + req := &bschemas.BifrostChatRequest{Provider: bschemas.OpenAI, Input: msgs} + rep := &components.Report{} + if _, err := p.Offload(req, rep, c); err != nil { + t.Fatalf("Offload: %v", err) + } + return req.Input, rep +} + +func newPin(t *testing.T) *Prefixpin { + t.Helper() + comp, err := newPrefixpin(nil) + if err != nil { + t.Fatal(err) + } + return comp.(*Prefixpin) +} + +// The core behaviour: a churning early block gets pinned back to its first +// rendering, so the prefix hashes identically and the provider cache still reads. +func TestPinsChurningEarlyBlock(t *testing.T) { + p, c := newPin(t), pinCtx() + first := scratchpad("FIRST") + + // turn 1: baseline recorded, nothing changed + got, rep := turn(t, p, c, []bschemas.ChatMessage{ + msg(bschemas.ChatMessageRoleUser, first), + msg(bschemas.ChatMessageRoleUser, "go"), + }) + if !rep.Skipped || schema.MessageText(got[0]) != first { + t.Fatal("turn 1 should only baseline, not modify") + } + + // turn 2: churn observed once -- still below RepeatThreshold, so no pin yet + second := scratchpad("SECOND") + got, _ = turn(t, p, c, []bschemas.ChatMessage{ + msg(bschemas.ChatMessageRoleUser, second), + msg(bschemas.ChatMessageRoleUser, "go"), + }) + if schema.MessageText(got[0]) != second { + t.Fatal("pinned on the first edit; a one-off edit must not be pinned") + } + + // turn 3: churn is now an established pattern -> pin back to the first text + third := scratchpad("THIRD") + got, rep = turn(t, p, c, []bschemas.ChatMessage{ + msg(bschemas.ChatMessageRoleUser, third), + msg(bschemas.ChatMessageRoleUser, "go"), + }) + if rep.Skipped { + t.Fatal("did not act on an established churn pattern -- this is the case that pays") + } + if schema.MessageText(got[0]) != first { + t.Fatalf("did not pin to the first rendering; got %.60q", schema.MessageText(got[0])) + } +} + +// An append-only conversation is the common case (100% of claude-code's turns, +// 98% of Bob's). Prefixpin must be completely inert there. +func TestInertWhenPrefixStable(t *testing.T) { + p, c := newPin(t), pinCtx() + head := scratchpad("ONLY") + for turnNo := 0; turnNo < 5; turnNo++ { + msgs := []bschemas.ChatMessage{msg(bschemas.ChatMessageRoleUser, head)} + for k := 0; k <= turnNo; k++ { + msgs = append(msgs, msg(bschemas.ChatMessageRoleAssistant, "step")) + } + got, rep := turn(t, p, c, msgs) + if !rep.Skipped { + t.Fatalf("turn %d: acted on a stable prefix", turnNo) + } + if schema.MessageText(got[0]) != head { + t.Fatalf("turn %d: mutated a stable message", turnNo) + } + } +} + +// If a DIFFERENT message occupies the slot (the transcript was restructured rather +// than edited in place), pinning would substitute unrelated content. It must +// re-baseline instead. +func TestDoesNotPinUnrelatedContent(t *testing.T) { + p, c := newPin(t), pinCtx() + a := strings.Repeat("alpha content here.\n", 300) + b := strings.Repeat("completely different beta text.\n", 300) + for i := 0; i < 4; i++ { + text := a + if i > 0 { + text = b + } + got, _ := turn(t, p, c, []bschemas.ChatMessage{ + msg(bschemas.ChatMessageRoleUser, text), + msg(bschemas.ChatMessageRoleUser, "go"), + }) + if i > 0 && schema.MessageText(got[0]) != b { + t.Fatalf("turn %d: substituted unrelated content", i) + } + } +} + +// The newest message is what the agent is actively reasoning about; pinning it +// would feed the model stale work. It must never be touched. +func TestNeverPinsNewestMessage(t *testing.T) { + p, c := newPin(t), pinCtx() + for i, name := range []string{"A", "B", "C", "D"} { + msgs := []bschemas.ChatMessage{msg(bschemas.ChatMessageRoleUser, scratchpad(name))} + got, _ := turn(t, p, c, msgs) + if schema.MessageText(got[len(got)-1]) != scratchpad(name) { + t.Fatalf("turn %d: modified the newest message", i) + } + } +} + +// Small blocks are not worth the behavioural risk of showing stale text. +func TestSkipsSmallBlocks(t *testing.T) { + p, c := newPin(t), pinCtx() + for i, name := range []string{"A", "B", "C", "D"} { + got, _ := turn(t, p, c, []bschemas.ChatMessage{ + msg(bschemas.ChatMessageRoleUser, "tiny "+name), + msg(bschemas.ChatMessageRoleUser, "go"), + }) + if schema.MessageText(got[0]) != "tiny "+name { + t.Fatalf("turn %d: pinned a block below MinTokens", i) + } + } +} + +// Deep messages must be left alone: only structurally-early slots are pinnable. +func TestOnlyPinsEarlyIndices(t *testing.T) { + p, c := newPin(t), pinCtx() + deep := 8 // beyond the default MaxPinIndex of 4 + for i, name := range []string{"A", "B", "C", "D"} { + msgs := make([]bschemas.ChatMessage, 0, deep+2) + for k := 0; k < deep; k++ { + msgs = append(msgs, msg(bschemas.ChatMessageRoleUser, "filler message body")) + } + msgs = append(msgs, msg(bschemas.ChatMessageRoleUser, scratchpad(name))) + msgs = append(msgs, msg(bschemas.ChatMessageRoleUser, "go")) + got, _ := turn(t, p, c, msgs) + if schema.MessageText(got[deep]) != scratchpad(name) { + t.Fatalf("turn %d: pinned a message past MaxPinIndex", i) + } + } +} + +// Lossy by design, so the withheld text MUST be recoverable: the current rendering +// is stashed under the returned key. +func TestStashesOriginalForExpand(t *testing.T) { + p, c := newPin(t), pinCtx() + texts := []string{scratchpad("ONE"), scratchpad("TWO"), scratchpad("THREE")} + var keys []string + for _, tx := range texts { + req := &bschemas.BifrostChatRequest{Provider: bschemas.OpenAI, Input: []bschemas.ChatMessage{ + msg(bschemas.ChatMessageRoleUser, tx), + msg(bschemas.ChatMessageRoleUser, "go"), + }} + k, err := p.Offload(req, &components.Report{}, c) + if err != nil { + t.Fatal(err) + } + keys = append(keys, k...) + } + if len(keys) == 0 { + t.Fatal("no rewind keys returned; the withheld text would be unrecoverable") + } + got, ok := c.Store.Get(keys[len(keys)-1]) + if !ok || string(got) != texts[len(texts)-1] { + t.Fatal("stashed original does not match what was withheld") + } +} + +// REGRESSION, from the real trace. The churning block re-rendered its iteration +// counter in ~20 scattered places: 152 changed chars out of 6,024 (98.5% identical), +// but with edits near BOTH ends. A prefix+suffix overlap measure scored that 0.075 +// and the guard rejected the exact case this component exists for. similarity must +// score it high. +func TestSimilarityToleratesScatteredEdits(t *testing.T) { + mk := func(word, num string) string { + var b strings.Builder + b.WriteString("\nInstruction seen for the " + word + " time.\n") + for i := 0; i < 60; i++ { + b.WriteString("stable reasoning line " + strconv.Itoa(i) + "\n") + } + // counter re-rendered in several forms, spread through the block + for i := 0; i < 8; i++ { + b.WriteString("iteration " + num + " of the " + strings.ToLower(word) + " pass\n") + b.WriteString("more stable content line " + strconv.Itoa(i) + "\n") + } + b.WriteString("[DONE] Generate " + strings.ToLower(word) + " state snapshot\n") + return b.String() + } + a, b := mk("THIRTY-SECOND", "32"), mk("THIRTY-THIRD", "33") + if a == b { + t.Fatal("fixture is not actually different") + } + s := similarity(a, b) + if s < 0.80 { + t.Fatalf("scattered-edit rewrite scored %.3f, below the 0.80 gate — the "+ + "component would skip the case it exists to fix", s) + } +} + +// similarity must actually separate the two cases it gates on. +func TestSimilarityDiscriminates(t *testing.T) { + a := scratchpad("THIRTY-SECOND") + b := scratchpad("THIRTY-THIRD") + if s := similarity(a, b); s < 0.9 { + t.Fatalf("counter-only edit scored %.3f; should be near 1", s) + } + c := strings.Repeat("totally unrelated text.\n", 300) + if s := similarity(a, c); s > 0.2 { + t.Fatalf("unrelated content scored %.3f; should be near 0", s) + } + if similarity("", "x") != 0 || similarity("x", "x") != 1 { + t.Fatal("degenerate cases wrong") + } +} + +// No store must not panic; it simply cannot track slots. +func TestNoStoreIsSafe(t *testing.T) { + p := newPin(t) + req := &bschemas.BifrostChatRequest{Provider: bschemas.OpenAI, Input: []bschemas.ChatMessage{ + msg(bschemas.ChatMessageRoleUser, scratchpad("A")), + msg(bschemas.ChatMessageRoleUser, "go"), + }} + rep := &components.Report{} + if _, err := p.Offload(req, rep, &components.Ctx{}); err != nil { + t.Fatal(err) + } + if !rep.Skipped { + t.Fatal("should skip without a store") + } +} diff --git a/docs/components/prefixpin.md b/docs/components/prefixpin.md new file mode 100644 index 0000000..2be9f8a --- /dev/null +++ b/docs/components/prefixpin.md @@ -0,0 +1,85 @@ +# prefixpin + +!!! warning "Offload — lossy, reversible" + Re-sends the **first** rendering of an early message that the agent rewrites in + place every turn, so the cached prefix stays byte-identical. The original is + stashed, so [`expand`](../how-to/recover-context.md) can recover it. + +## The problem it solves + +[`cacheinject`](cacheinject.md) optimises *where* the cache boundary goes. prefixpin +fixes the case where **no boundary can help**: an agent that mutates an early message +on every turn. + +Providers hash the request prefix cumulatively. A single changed character at message +index 1 makes every token above it unmatchable — there is no breakpoint position whose +hash excludes an earlier block. + +Measured across 1,955 real Bob requests on SWE-bench: + +| | requests | uncached input | +|---|--:|--:| +| append-only turns (cached fine) | 98.0% | 28.2% | +| **one early message mutated** | **2.0%** | **71.8%** (5,796,220 tokens) | + +On one task the agent re-emitted a running ``/`` at index +1, re-rendering an iteration counter in ~20 places (`"THIRTY-SECOND"` → +`"THIRTY-THIRD"`, `"32"` → `"33"`): **152 changed characters out of 6,024 — 98.5% +identical content** — sitting ~1,374 tokens into a 181k-token prefix. Only 0.76% of the +prefix survived; the cache-hit rate collapsed from 98% to 5.7%. + +The economics are lopsided. A cache read costs 0.1× base input and uncached input +costs 1.0×, so a mutation below the boundary makes every token above it **ten times** +more expensive. Placement tuning only ever moves tokens between read (0.1×) and +cache-write (1.25×) — which is why this is worth ~31% of Bob's input cost while every +placement change measured ~0%. + +## How it works + +For each early message, remember the first text seen for its `(index, role)` slot in +this session. On a later turn, if that slot's text **changed** but is still +recognisably the same content (same shape, high similarity), re-send the **first** +rendering so the prefix stays byte-identical. + +## Lossiness + +Deliberately an Offload, not a Reformat. The model sees the pinned (older) text rather +than what the agent just wrote, so information **is** withheld: a counter reads stale. +That is a real behavioural change, and the reason for the guards. + +## Guards + +Each closes a way this could do harm: + +- only messages at index `< max_pin_index` — an early, structural slot, never the + working tail the agent is actively reasoning about; +- only when similarity `>= min_similarity` — a genuinely rewritten-in-place block, not + a different message that happens to occupy the slot; +- only after the slot has churned `repeat_threshold` times, so a one-off edit is never + pinned — only a per-turn churn *pattern*; +- never on the newest message, never on tool results. + +## Configuration + +```yaml +pipeline: [prefixpin, cacheinject] +components: + prefixpin: + max_pin_index: 4 # 0 disables + min_similarity: 0.80 + repeat_threshold: 2 + min_tokens: 200 +``` + +| key | default | meaning | +|---|--:|---| +| `max_pin_index` | `4` | bounds pinning to structurally-early messages; `0` disables | +| `min_similarity` | `0.80` | character-level overlap required to treat a changed slot as the same content rewritten in place | +| `repeat_threshold` | `2` | how many times a slot must churn before pinning starts | +| `min_tokens` | `200` | skips slots too small to be worth the behavioural risk | + +## When to enable it + +Enabled for every provider: the failure mode is provider-independent. It bites hardest +on **implicit-cache backends** (Gemini/Bob, OpenAI) where there is no `cache_control` +to place and prefix stability is the *only* available lever. diff --git a/mkdocs.yml b/mkdocs.yml index dd77d02..4e50694 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -128,6 +128,7 @@ nav: - extract_llm: components/extract_llm.md - smartcrush: components/smartcrush.md - mask: components/mask.md + - prefixpin: components/prefixpin.md - summarize: components/summarize.md - The DSL filter engine: components/dsl.md - How-to Guides: