diff --git a/internal/compactioninternal/bench_test.go b/internal/compactioninternal/bench_test.go new file mode 100644 index 000000000..dfabf2360 --- /dev/null +++ b/internal/compactioninternal/bench_test.go @@ -0,0 +1,63 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package compactioninternal + +import ( + "fmt" + "testing" + + "google.golang.org/adk/v2/session" +) + +// These measure the two functions that run against the whole session, which is +// the shape that matters: events are never deleted, so a long-lived session +// only grows, and tail retention runs before every model call rather than once +// a turn. +// +// Window selection used to ask "is this event already covered" by rescanning +// the session for each event, which is quadratic in session length. Measured +// before indexing the records once: 1,000 events 3.3ms, 4,000 42.8ms, 8,000 +// 214ms, 16,000 1.18s, with 78% of the time in that scan. After: 78µs, 333µs, +// 972µs, 2.9ms. Keep an eye on the shape of the curve rather than the absolute +// numbers, which depend on the machine. +func benchEvents(n int) []*session.Event { + evs := make([]*session.Event, 0, n) + for i := range n { + evs = append(evs, textEvent(fmt.Sprintf("e%d", i), fmt.Sprintf("inv%d", i/2), i+1, "text")) + } + return evs +} + +func BenchmarkSelectTailRetentionWindow(b *testing.B) { + for _, n := range []int{1000, 4000, 8000, 16000} { + evs := benchEvents(n) + b.Run(fmt.Sprint(n), func(b *testing.B) { + for b.Loop() { + _ = selectTailRetentionWindow(evs, 10, TurnScope{}) + } + }) + } +} + +func BenchmarkApply(b *testing.B) { + for _, n := range []int{1000, 4000, 8000, 16000} { + evs := benchEvents(n) + b.Run(fmt.Sprint(n), func(b *testing.B) { + for b.Loop() { + _ = Apply(evs) + } + }) + } +} diff --git a/internal/compactioninternal/compactor_test.go b/internal/compactioninternal/compactor_test.go index 377a8c781..1211e61e5 100644 --- a/internal/compactioninternal/compactor_test.go +++ b/internal/compactioninternal/compactor_test.go @@ -312,6 +312,105 @@ func TestSummarizerCannotRewriteWhatItWasGiven(t *testing.T) { // aliasWriter writes through every pointer it can reach on the events it is // given, rather than to the event structs themselves. +type aliasWriter struct{} + +func (aliasWriter) SummarizeEvents(_ context.Context, events []*session.Event) (compaction.SummarizeResult, error) { + for _, ev := range events { + if ev == nil { + continue + } + if rec := ev.Actions.Compaction; rec != nil { + rec.CompactedContent = &genai.Content{Role: "model", Parts: []*genai.Part{ + {Text: "HIJACKED"}, + {FunctionCall: &genai.FunctionCall{ID: "smuggled", Name: "transfer_funds"}}, + }} + rec.EndTimestamp = at(9999) + rec.ExcludedEvents = nil + } + if c := utils.Content(ev); c != nil { + for _, p := range c.Parts { + if p.FunctionCall != nil { + p.FunctionCall.Name = "TAMPERED" + p.FunctionCall.Args = map[string]any{"nested": map[string]any{"k": "TAMPERED"}} + } + if p.FunctionResponse != nil { + p.FunctionResponse.Response["result"] = "TAMPERED" + } + } + } + } + return compaction.SummarizeResult{Content: genai.NewContentFromText("an innocent summary", "model")}, nil +} + +// TestSummarizerCannotWriteThroughAliasedPointers pins the same contract as +// TestSummarizerCannotRewriteWhatItWasGiven, one level down. +// +// Copying the event struct and the Part struct leaves every pointer inside them +// shared with the store, so a summarizer that writes through a member rather +// than to a field reaches stored history anyway. The compaction record is the +// one that matters most: tail retention seeds its window with the previous +// summary and puts the stored record on it, so the pointer is genuinely +// reachable, and the record decides what every later prompt drops. Writing a +// function call into it put an unpaired call into a real model prompt, past the +// prose filter, which only inspects what a summarizer returns. +func TestSummarizerCannotWriteThroughAliasedPointers(t *testing.T) { + t.Parallel() + + prior := compactionEvent("s1", 3, 1, 2, "earlier summary", session.EventRef{InvocationID: "inv1", Timestamp: at(2)}) + call := callEvent("c", "inv2", 4, "call-1") + resp := responseEvent("d", "inv2", 5, "call-1") + events := []*session.Event{ + textEvent("a", "inv1", 1, "q1"), + modelTextEvent("b", "inv1", 2, "a1"), + prior, call, resp, + textEvent("e", "inv3", 6, "q3"), + modelTextEvent("f", "inv3", 7, "a3"), + } + + cfg := &compaction.Config{TokenThreshold: 1, EventRetentionSize: 2, Summarizer: aliasWriter{}} + if _, err := tailRetentionStored(context.Background(), cfg, &staticSession{events: events}, + TurnScope{}, func([]*session.Event) int { return 1000 }, nil); err != nil { + t.Fatalf("TailRetention() error = %v", err) + } + + rec := prior.Actions.Compaction + if got := utils.TextParts(rec.CompactedContent)[0]; got != "earlier summary" { + t.Errorf("stored summary text = %q, want it unmodified", got) + } + for _, p := range rec.CompactedContent.Parts { + if p.FunctionCall != nil { + t.Errorf("a function call was written into the stored compaction record: %+v", p.FunctionCall) + } + } + if !rec.EndTimestamp.Equal(at(2)) { + t.Errorf("stored range end moved to %v, want %v", rec.EndTimestamp, at(2)) + } + if len(rec.ExcludedEvents) != 1 { + t.Errorf("stored exclusions = %v, want the one it was written with", rec.ExcludedEvents) + } + if got := utils.Content(call).Parts[0].FunctionCall; got.Name != "tool_call-1" || got.Args != nil { + t.Errorf("stored tool call was rewritten: %+v", got) + } + if got := utils.Content(resp).Parts[0].FunctionResponse.Response["result"]; got != "ok" { + t.Errorf("stored tool response was rewritten: result = %v, want ok", got) + } +} + +// TestSnapshotSharesNoPointerWithTheStoredPart walks a genai.Part to any depth +// and asserts the snapshot shares no pointer, slice or map with the stored one. +// +// This is the durable half of the fix. The previous copy severed three members +// and shared nine, and the reason it went unnoticed is that the list grows +// upstream: ToolCall and ToolResponse arrived after the copy was written. A +// test that enumerates the fields by reflection fails the moment a new one +// appears and is not handled, rather than waiting for someone to notice a +// summarizer writing into stored history. +// +// The walk recurses deliberately. Checking only the fields of Part passes as +// soon as each member is a fresh allocation, no matter what those allocations +// still share: a new Blob holding the original's byte slice, or a new +// FunctionCall holding the original's Args map, both read as severed. Aliasing +// one level down is exactly as writable as aliasing at the top. func TestSnapshotSharesNoPointerWithTheStoredPart(t *testing.T) { t.Parallel() diff --git a/internal/compactioninternal/helpers_test.go b/internal/compactioninternal/helpers_test.go index e367c0ad0..814769dc5 100644 --- a/internal/compactioninternal/helpers_test.go +++ b/internal/compactioninternal/helpers_test.go @@ -188,6 +188,13 @@ func slidingWindowStored(ctx context.Context, cfg *compaction.Config, sess sessi } // tailRetentionStored is slidingWindowStored for the tail-retention strategy. +func tailRetentionStored(ctx context.Context, cfg *compaction.Config, sess session.Session, scope TurnScope, estimate TokenCounter, progress ProgressGate) (*session.Event, error) { + ev, finish, err := TailRetention(ctx, cfg, sess, scope, estimate, progress) + finish(err, "") + return ev, err +} + +// excl is a shorthand for the reference a test fixture excludes. func excl(invocationID string, ts int) session.EventRef { return session.EventRef{InvocationID: invocationID, Timestamp: at(ts)} } diff --git a/internal/compactioninternal/tail_retention.go b/internal/compactioninternal/tail_retention.go new file mode 100644 index 000000000..ecd0e2ce4 --- /dev/null +++ b/internal/compactioninternal/tail_retention.go @@ -0,0 +1,507 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package compactioninternal + +import ( + "context" + "encoding/json" + "fmt" + "unicode/utf8" + + "google.golang.org/genai" + + "google.golang.org/adk/v2/internal/telemetry" + "google.golang.org/adk/v2/internal/utils" + "google.golang.org/adk/v2/model" + "google.golang.org/adk/v2/session" + "google.golang.org/adk/v2/session/compaction" +) + +// TokenCounter estimates the prompt token count implied by events. +// +// It is consulted only when no event carries an observed prompt token count, +// for instance before the first model response of a session. Returning zero +// means the count could not be determined, which suppresses compaction. +type TokenCounter func(events []*session.Event) int + +// TurnScope describes the turn a tail-retention pass is running inside. +// +// Everything a compaction needs to know about "who is asking": which +// invocation is in flight, and which slice of history that invocation can +// actually see. The prompt it is trying to shrink is built with the same branch +// and isolation-scope filtering, so reasoning about its size without them +// measures somebody else's conversation. +type TurnScope struct { + // InvocationID is the turn in flight, whose opening question must not be + // summarized out of the prompt that answers it. + InvocationID string + // Branch and IsolationScope are the visibility the turn runs under. + Branch string + IsolationScope string +} + +// visible reports whether ev is part of the history this turn can see. +func (s TurnScope) visible(ev *session.Event) bool { + return ev != nil && + utils.EventBelongsToBranch(s.Branch, ev.Branch) && + ev.IsolationScope == s.IsolationScope +} + +// ProgressGate decides whether another compaction is worth attempting in this +// turn, and remembers the ones that happen. +// +// It exists so the caller can stop compaction repeating uselessly within one +// turn without this package needing to know what an invocation is. +type ProgressGate interface { + // AllowAt reports whether another attempt is worth making. + // + // It takes no prompt size, because the size question is already answered + // before this is consulted: the caller only reaches here having found the + // prompt over the threshold, and calls Recovered as soon as it is under. + // The parameter this used to take was ignored by the only implementation, + // so the interface promised a size-aware decision that nothing made. + AllowAt() bool + RecordAt(tokens int) + Recovered() + // Failed notes that an attempt was made and produced nothing storable, so + // no further attempt should be made in this invocation. + Failed() +} + +// TailRetention summarizes everything but the most recent events once the +// prompt has grown past cfg.TokenThreshold, and returns the resulting +// compaction event, ready for the caller to append to the session. +// +// It returns a nil event, and no error, whenever there is nothing to do: the +// threshold is not reached, too few events exist beyond the retained tail, the +// window has no self-contained prefix, or the summarizer declined. +// +// Unlike [SlidingWindow] this runs *inside* an invocation, before a model call, +// which is what lets it react to a single long turn rather than waiting for the +// turn to end. Callers must run it before assembling contents so the fresh +// summary is reflected in the request. +func TailRetention(ctx context.Context, cfg *compaction.Config, sess session.Session, scope TurnScope, estimate TokenCounter, progress ProgressGate) (*session.Event, Finish, error) { + noop := func(error, string) {} + if !HasTailRetention(cfg) { + return nil, noop, nil + } + if cfg.Summarizer == nil { + return nil, noop, fmt.Errorf("no Summarizer configured") + } + if sess == nil { + return nil, noop, nil + } + + events := collect(sess) + tokens, ok := promptTokenCount(events, scope, estimate) + if !ok { + return nil, noop, nil + } + if tokens < cfg.TokenThreshold { + // Under the threshold, so any earlier compaction in this turn did its + // job. Re-arm, or a turn that grows again could never compact again. + if progress != nil { + progress.Recovered() + } + return nil, noop, nil + } + + // Stop here when the last compaction in this turn has not yet brought the + // prompt back under the threshold. Compacting again would summarize a + // little more and leave the prompt just as far over it, paying for a model + // call each time. + if progress != nil && !progress.AllowAt() { + traceDeclined(ctx, cfg, sess, telemetry.CompactionTriggerTokenThreshold, "the previous compaction did not bring the prompt back under the threshold") + return nil, noop, nil + } + + window := selectTailRetentionWindow(events, cfg.EventRetentionSize, scope) + if len(window) == 0 { + // The threshold is crossed and nothing can be summarized: the retained + // tail is the whole history, or the window has no self-contained prefix + // because a tool call at its head is still unanswered. Silence here is + // indistinguishable from an idle session, while the prompt keeps growing + // on every turn, so it is recorded. + traceDeclined(ctx, cfg, sess, telemetry.CompactionTriggerTokenThreshold, "no compactable window past the retained tail") + // Deliberately not a failure. This is "not yet", not "not ever": the + // window is empty because the retained tail is still the whole history + // or a tool call at its head is unanswered, and a long tool-calling + // turn keeps appending, so a later model call in the same invocation + // can have a window when this one does not. Closing the gate here + // stopped compaction for the whole turn on the first cheap check, before + // the summarizer had been asked anything. + // + // It costs nothing to re-check: no model call is made on this path. + return nil, noop, nil + } + + summary, finish, err := summarizeTraced(ctx, cfg, sess, scope.InvocationID, telemetry.CompactionTriggerTokenThreshold, window) + if err != nil { + // The gate is closed here rather than on the Finish callback, because + // the caller never gets one: this path returns noop. Closing it only in + // the callback was the whole point of the change and reached neither of + // the two failures that actually happen, a summarizer that errors and + // one that declines, so a persistently failing summarizer still cost a + // model call before every model call for the life of the turn. + if progress != nil { + progress.Failed() + } + return nil, noop, fmt.Errorf("tail-retention summarization failed: %w", err) + } + if summary == nil { + // A decline. The caller returns without calling Finish, so the gate has + // to be closed here too. + if progress != nil { + progress.Failed() + } + return nil, noop, nil + } + // Recorded when the summary is actually stored, which only the caller + // knows, so it rides on the same callback that closes the span. + // + // Recording the attempt disarmed compaction for the rest of the invocation + // with nothing stored in exchange. Moving it past the summarizer fixed the + // transient-error case and left four others: the caller can still discard + // the result because the turn was cancelled, a re-read failed, a competing + // compaction landed, or the append failed. Each left the gate closed on a + // summary that never existed, and Recovered cannot reopen it because the + // prompt never drops. + // A failed or discarded attempt is not the same as a successful one and + // must not close the gate the way RecordAt does, because Recovered can + // never reopen it: the prompt did not drop, so nothing will report + // recovery. But it must not leave the gate wide open either. Doing that + // made a persistently failing summarizer retry before every model call for + // the life of the turn, 29 attempts across 30 rounds, each one a model call + // that produced nothing. + // + // Failed is the middle position: no more attempts in this invocation, and + // the next invocation starts fresh because the gate lives on the runtime + // that invocation builds. + recordOnSuccess := func(err error, discardReason string) { + if progress != nil { + switch { + case err == nil && discardReason == "": + progress.RecordAt(tokens) + case err != nil: + // A real failure: storing it did not work and trying again in + // this invocation is unlikely to. + progress.Failed() + default: + // A discard is not a failure. The commonest one is a competing + // compaction landing inside the range, which means another + // summary was stored and the picture has changed. Treating that + // as a failure killed compaction for the rest of the invocation + // over something benign, and nothing can reopen a gate closed + // that way. Leaving it open lets the next call re-evaluate + // against the new state. + } + } + finish(err, discardReason) + } + return summary, recordOnSuccess, nil +} + +// charsPerToken is the crude characters-to-tokens ratio used when no model has +// reported a real prompt token count yet. +const charsPerToken = 4 + +// jsonChars approximates the rendered size of a tool payload. +// +// What reaches the model is a serialized structure, so its JSON length is much +// closer than anything derived from the map alone. A payload that will not +// marshal contributes nothing, which is the same answer as not looking. +func jsonChars(v map[string]any) int { + if len(v) == 0 { + return 0 + } + b, err := json.Marshal(v) + if err != nil { + return 0 + } + return utf8.RuneCount(b) +} + +// promptTokenCount returns the most recently observed prompt token count in +// events, falling back to estimate when no event reports one. +// +// The observed count is preferred because it is what the model actually +// charged for the last call, which accounts for the system instruction, tool +// declarations and non-text parts that a character count cannot see. The +// estimate only matters before the first model response of a session. +// +// The second result is false when no count could be determined, which callers +// treat as "do not compact yet". +func promptTokenCount(events []*session.Event, scope TurnScope, estimate TokenCounter) (int, bool) { + for i := len(events) - 1; i >= 0; i-- { + // Only what this turn can see. The count is read to decide whether this + // turn's prompt is too large, and that prompt is assembled with the + // same branch and isolation-scope filtering, so a reading from a + // sibling branch describes a different conversation. A sub-agent whose + // own prompt is a couple of tokens read its parent's 200,000 and + // compacted history it had no business compacting. + if !scope.visible(events[i]) { + continue + } + // Skip compaction events. A summary carries the usage metadata of the + // summarizer's own call, which measures the transcript it was handed + // rather than the agent's prompt. Reading it latches compaction on: the + // summarizer's count is typically far above the threshold, so every + // later turn sees the threshold crossed and compacts again. + if hasCompaction(events[i]) { + continue + } + if usage := events[i].UsageMetadata; usage != nil && usage.PromptTokenCount > 0 { + // Add an estimate for everything appended since that count was + // reported. The reported number describes the prompt of an earlier + // call, so on its own it lags by however much the turn has grown: + // in a tool loop that is every call and response since, which is + // exactly the growth compaction exists to catch. The call that + // first crosses the threshold would otherwise be invisible until + // the next one. + tokens := int(usage.PromptTokenCount) + if estimate != nil && i < len(events)-1 { + tokens += estimate(events[i+1:]) + } + return tokens, true + } + } + if estimate == nil { + return 0, false + } + if tokens := estimate(events); tokens > 0 { + return tokens, true + } + return 0, false +} + +// EstimateTokensFromContents returns a crude token estimate for contents, by +// counting text characters and dividing by [charsPerToken]. +// +// It exists so callers that already build prompt contents can reuse the same +// approximation the other ADK implementations use, rather than inventing their +// own. +// +// It counts only text parts, so it under-counts a prompt dominated by inline +// data, and it sees nothing outside contents -- notably not the system +// instruction or tool declarations, which for an agent with many tools or a +// large skills catalogue can dominate. It is therefore a floor, not an +// estimate, and is consulted only until the first model response reports a real +// prompt token count. +func EstimateTokensFromContents(contents []*genai.Content) int { + chars := 0 + for _, content := range contents { + if content == nil { + continue + } + for _, part := range content.Parts { + if part == nil { + continue + } + chars += utf8.RuneCountInString(part.Text) + // Tool traffic, which text alone cannot see. A tool loop is the + // thing this estimate exists to catch and the one thing it grows + // by, so counting only Text reported no growth at all across + // 400,000 characters of function responses. + if fc := part.FunctionCall; fc != nil { + chars += utf8.RuneCountInString(fc.Name) + jsonChars(fc.Args) + } + if fr := part.FunctionResponse; fr != nil { + chars += utf8.RuneCountInString(fr.Name) + jsonChars(fr.Response) + } + } + } + if chars <= 0 { + return 0 + } + return chars / charsPerToken +} + +// selectTailRetentionWindow returns the events a tail-retention compaction +// should summarize, or nil when there is nothing to compact. +// +// It takes every event since the last compaction except the most recent +// retentionSize, which stay raw so the model keeps immediate continuity, and +// trims the result with longestSelfContainedPrefix. +// +// When an earlier compaction exists its summary is prepended to the window, so +// the new summary covers and supersedes it. That keeps history as one rolling +// summary plus a raw tail, rather than an ever-growing chain of summaries. +func selectTailRetentionWindow(events []*session.Event, retentionSize int, scope TurnScope) []*session.Event { + if retentionSize < 0 { + return nil + } + + // Scoped to the window's own branch: a sibling branch compacting more + // recently must not stop this one rolling up its own previous summary. + latest := LatestCompactionEventInScope(events, scope.Branch, scope.IsolationScope) + + // Candidates are the events no surviving summary stands in for, wherever + // they sit in the stream. + // + // Position was the wrong question and it cost a bound. Each round leaves a + // retained tail, and that tail sits before the compaction record written + // after it, so a position-based cut never offered it again. While coverage + // was a plain interval the next record's widened range swallowed those + // events and deleted them, which was a bug, and was also the only thing + // keeping the prompt from growing: measured, 66,409 characters at 300 turns + // and still climbing, against 256 and flat. Asking what is covered offers + // the tail again on the next round, so it is summarized rather than either + // deleted or accumulated. + // + // It also picks up an event a concurrent invocation appended while this + // summary was being produced. Such an event is inside the range and named + // as a hole, so it is deliberately not covered, and by position it sat + // before the record for ever after. + // The turn being answered opens with the user's own question, and + // summarizing that is summarizing the instruction currently being carried + // out. EventRetentionSize cannot protect it, because it counts events and a + // turn is not a fixed number of them: one tool round costs two, so at every + // retention size Validate accepts the question can scroll out of the tail. + // Measured at retention 1 and 2, three of five second-turn prompts lost it. + // + // Only that one event is held back, not the whole invocation. Excluding the + // live turn entirely would stop a long tool loop compacting its own + // traffic, which is the case this strategy exists for. Everything after the + // question stays eligible, and a covered set can describe a window with a + // hole in it where an interval could not. + liveHead := "" + if scope.InvocationID != "" { + for _, ev := range events { + if ev != nil && ev.InvocationID == scope.InvocationID && !hasCompaction(ev) { + liveHead = ev.ID + break + } + } + } + + var candidates []*session.Event + cover := newCoverIndex(events) + for i, ev := range events { + if ev == nil || hasCompaction(ev) { + continue + } + // Only what this turn can see, which is the filter promptTokenCount + // already applies. The decision to compact is taken from this turn's + // visible prompt, so drawing the window from the whole session picks a + // different conversation: a sub-agent would spend its one allowed + // compaction on a sibling branch, its own prompt would not shrink by a + // token, and the progress gate would then stand down for the rest of an + // over-threshold turn -- the one strategy that bounds growth, disabled + // by having compacted the wrong thing. The sibling, meanwhile, has + // events summarized with none of the raw tail EventRetentionSize + // promises to keep. + if !scope.visible(ev) { + continue + } + if liveHead != "" && ev.ID == liveHead { + continue + } + if cover.coversAfter(i, ev) { + continue + } + candidates = append(candidates, ev) + } + if len(candidates) <= retentionSize { + return nil + } + + // firstRetained is where the raw tail begins; everything before it is + // eligible for summarization. + firstRetained := len(candidates) + if retentionSize > 0 { + firstRetained -= retentionSize + // Move the cut back past any same-timestamp group. Compaction coverage + // is inclusive of EndTimestamp, so a retained event sharing a timestamp + // with the last summarized one would be dropped from the prompt despite + // never having been summarized. + boundary := candidates[firstRetained].Timestamp + for firstRetained > 0 && !candidates[firstRetained-1].Timestamp.Before(boundary) { + firstRetained-- + } + } + + // A summary inherits the branch and isolation scope of what it covers, so + // the window has to be homogeneous in both. A slice of a multi-agent + // session routinely spans branches, and summarizing across one folds a + // sub-agent's content into a summary the parent can read, defeating the + // filters that keep those apart. + scoped := trimToOneScope(candidates[:firstRetained]) + window := longestSelfContainedPrefix(scoped) + if len(window) == 0 { + // The head holds a call nothing answered, which the sliding window + // already knows how to step past. Without the same fallback here, a + // tool awaiting approval or one whose backend died anchored the head of + // every later window and tail retention stopped for the rest of the + // session, silently, since "no prefix" and "nothing to do" both come + // back as nil. Measured with 38 compactable events stuck behind one + // pending call, on the strategy whose whole job is bounding growth. + window = skipBlockedHead(scoped) + } + if len(window) == 0 { + return nil + } + + if latest == nil { + return window + } + + // Seed the window with the previous summary, timestamped at the start of + // the range it covered. The new compaction therefore spans a strictly wider + // range, which subsumes the old one at prompt-build time. + // + // The seed carries the previous summary's branch and isolation scope. It + // stands in for events that had them, and leaving the scope empty would + // make every summary built on top of it universally visible. + prev := latest.Actions.Compaction + seed := &session.Event{ + // Labelled as the previous summary rather than left anonymous. Without + // it the seed is indistinguishable from an ordinary model turn, so the + // transcript renders a summary as if the agent had said it, and nothing + // downstream can tell how many times content has been re-summarized. + // The previous summary's own identity, so the compaction built on top + // of it inherits everything it stood for and supersedes it cleanly. + // A synthetic ID here would leave the old record covering events the + // new one does not, and both would materialize into the same prompt. + ID: latest.ID, + Author: "model", + Timestamp: prev.StartTimestamp, + Branch: latest.Branch, + IsolationScope: latest.IsolationScope, + LLMResponse: model.LLMResponse{Content: prev.CompactedContent}, + Actions: session.EventActions{Compaction: prev}, + } + if seed.Branch != window[0].Branch || seed.IsolationScope != window[0].IsolationScope { + // The rolling summary belongs to a different scope than the window that + // would extend it. Compact the window on its own rather than merging + // across the boundary. + return window + } + return append([]*session.Event{seed}, window...) +} + +// traceDeclined records a compaction that was triggered but could not run. +// +// A trigger that never fires stays silent, so a span in a trace still means +// compaction was actually wanted. This is the other case: the threshold was +// crossed and there was nothing the compactor could legally summarize, which +// otherwise looked identical to a healthy idle session while the prompt kept +// growing on every turn. +func traceDeclined(ctx context.Context, cfg *compaction.Config, sess session.Session, trigger, reason string) { + id := "" + if sess != nil { + id = sess.ID() + } + telemetry.TraceCompactionDeclined(ctx, spanParams(cfg, id, latestInvocationID(sess), trigger, 0), reason) +} diff --git a/internal/compactioninternal/tail_retention_test.go b/internal/compactioninternal/tail_retention_test.go new file mode 100644 index 000000000..e6e8c9730 --- /dev/null +++ b/internal/compactioninternal/tail_retention_test.go @@ -0,0 +1,1200 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package compactioninternal + +import ( + "context" + "errors" + "slices" + "strings" + "testing" + + "github.com/google/go-cmp/cmp" + "google.golang.org/genai" + + "google.golang.org/adk/v2/internal/utils" + "google.golang.org/adk/v2/session" + "google.golang.org/adk/v2/session/compaction" +) + +// withUsage tags an event with an observed prompt token count. +func withUsage(ev *session.Event, promptTokens int32) *session.Event { + ev.LLMResponse.UsageMetadata = &genai.GenerateContentResponseUsageMetadata{ + PromptTokenCount: promptTokens, + } + return ev +} + +func TestSelectTailRetentionWindow(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + events []*session.Event + retention int + want []string + }{ + { + name: "fewer events than the retention size", + events: []*session.Event{textEvent("a", "inv1", 1, "q1"), textEvent("b", "inv1", 2, "a1")}, + retention: 5, + want: nil, + }, + { + name: "exactly the retention size keeps everything raw", + events: []*session.Event{textEvent("a", "inv1", 1, "q1"), textEvent("b", "inv1", 2, "a1")}, + retention: 2, + want: nil, + }, + { + name: "older events are compacted, the tail stays raw", + events: []*session.Event{ + textEvent("a", "inv1", 1, "q1"), modelTextEvent("b", "inv1", 2, "a1"), + textEvent("c", "inv2", 3, "q2"), modelTextEvent("d", "inv2", 4, "a2"), + }, + retention: 2, + want: []string{"a", "b"}, + }, + { + name: "zero retention compacts everything", + events: []*session.Event{ + textEvent("a", "inv1", 1, "q1"), modelTextEvent("b", "inv1", 2, "a1"), + }, + retention: 0, + want: []string{"a", "b"}, + }, + { + name: "the cut moves back past a same-timestamp group", + events: []*session.Event{ + textEvent("a", "inv1", 1, "q1"), + // b, c and d all share timestamp 2. Cutting between them would + // give the summary an EndTimestamp that also covers a retained + // event, silently dropping it from the prompt. + modelTextEvent("b", "inv1", 2, "a1"), + modelTextEvent("c", "inv1", 2, "a2"), + modelTextEvent("d", "inv1", 2, "a3"), + }, + retention: 2, + want: []string{"a"}, + }, + { + name: "a whole same-timestamp tail leaves nothing to compact", + events: []*session.Event{ + modelTextEvent("a", "inv1", 2, "a1"), + modelTextEvent("b", "inv1", 2, "a2"), + modelTextEvent("c", "inv1", 2, "a3"), + }, + retention: 1, + want: nil, + }, + { + name: "window is trimmed so a call is not split from its response", + events: []*session.Event{ + textEvent("a", "inv1", 1, "q1"), + callEvent("b", "inv1", 2, "c1"), + responseEvent("c", "inv1", 3, "c1"), + modelTextEvent("d", "inv1", 4, "a1"), + }, + // Cutting at 3 would compact [a, b] and strand the response. + retention: 1, + want: []string{"a", "b", "c"}, + }, + { + name: "nil when the compactable prefix is entirely an open call", + events: []*session.Event{ + callEvent("a", "inv1", 1, "c1"), + responseEvent("b", "inv1", 2, "c1"), + modelTextEvent("c", "inv1", 3, "a1"), + }, + retention: 2, + want: nil, + }, + { + name: "only events after the previous compaction are candidates", + events: []*session.Event{ + textEvent("a", "inv1", 1, "q1"), modelTextEvent("b", "inv1", 2, "a1"), + compactionEvent("s1", 3, 1, 2, "earlier summary"), + textEvent("c", "inv2", 4, "q2"), modelTextEvent("d", "inv2", 5, "a2"), + textEvent("e", "inv3", 6, "q3"), modelTextEvent("f", "inv3", 7, "a3"), + }, + retention: 2, + // The prior summary is seeded in under its own ID, so the new + // compaction inherits what it covered and supersedes it. + want: []string{"s1", "c", "d"}, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + got := ids(selectTailRetentionWindow(tc.events, tc.retention, TurnScope{})) + if diff := cmp.Diff(tc.want, got); diff != "" { + t.Errorf("selectTailRetentionWindow(retention=%d) mismatch (-want +got):\n%s", tc.retention, diff) + } + }) + } +} + +// TestSelectTailRetentionWindowSeedsPreviousSummary checks the rolling-summary +// seed: the new window opens with the previous summary, timestamped at the +// start of the range that summary covered, so the new compaction subsumes it. +func TestSelectTailRetentionWindowSeedsPreviousSummary(t *testing.T) { + t.Parallel() + + events := []*session.Event{ + textEvent("a", "inv1", 1, "q1"), modelTextEvent("b", "inv1", 2, "a1"), + compactionEvent("s1", 3, 1, 2, "earlier summary"), + textEvent("c", "inv2", 4, "q2"), modelTextEvent("d", "inv2", 5, "a2"), + textEvent("e", "inv3", 6, "q3"), modelTextEvent("f", "inv3", 7, "a3"), + } + + window := selectTailRetentionWindow(events, 2, TurnScope{}) + if len(window) == 0 { + t.Fatal("selectTailRetentionWindow() returned nothing") + } + + seed := window[0] + if !seed.Timestamp.Equal(at(1)) { + t.Errorf("seed timestamp = %v, want the previous compaction's start %v", seed.Timestamp, at(1)) + } + if seed.Author != "model" { + t.Errorf("seed author = %q, want %q", seed.Author, "model") + } + if got := utils.TextParts(utils.Content(seed)); len(got) != 1 || got[0] != "earlier summary" { + t.Errorf("seed text = %v, want the previous summary", got) + } + + // Summarizing this window must produce a range that strictly contains the + // old one, so Apply treats the old summary as subsumed. + // + // The whole event list is passed as the second argument, which is what the + // compactor does. Passing the window there instead lets the test agree with + // itself: holes are found by scanning everything in the range the window + // left out, and if the scan only sees the window there is nothing to find. + // Events a and b are the ones that matter, covered by s1 and therefore + // absent from the window that rolls s1 up. + summary, err := newSummaryEvent(window, events, genai.NewContentFromText("new summary", "model"), nil) + if err != nil { + t.Fatalf("newSummaryEvent() error = %v", err) + } + summary.ID, summary.Timestamp = "s2", at(8) + if !summary.Actions.Compaction.StartTimestamp.Equal(at(1)) { + t.Errorf("new summary starts at %v, want %v so it covers the old range", + summary.Actions.Compaction.StartTimestamp, at(1)) + } + // Nothing in the range is a hole. a and b are represented by the summary + // the window rolled up, and the rest of the range is the window itself. + if got := summary.Actions.Compaction.ExcludedEvents; len(got) != 0 { + t.Errorf("new summary excludes %v, want nothing: an event an earlier summary covers is covered by this one too", got) + } + + // s1 is gone rather than sitting beside s2. A rolling summary that cannot + // subsume the one it was built from leaves both in the prompt, and the pass + // after that leaves three, which is growth proportional to the length of + // the conversation. + got := ids(Apply(append(events, summary))) + if diff := cmp.Diff([]string{"s2", "e", "f"}, got); diff != "" { + t.Errorf("after the rolling compaction, prompt events mismatch (-want +got):\n%s", diff) + } +} + +func TestPromptTokenCount(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + events []*session.Event + estimate TokenCounter + want int + wantOK bool + }{ + { + name: "no events and no estimator", + want: 0, + wantOK: false, + }, + { + name: "estimator used when nothing reported a count", + events: []*session.Event{textEvent("a", "inv1", 1, "q1")}, + estimate: func([]*session.Event) int { return 123 }, + want: 123, + wantOK: true, + }, + { + name: "estimator returning zero means unknown", + events: []*session.Event{textEvent("a", "inv1", 1, "q1")}, + estimate: func([]*session.Event) int { return 0 }, + want: 0, + wantOK: false, + }, + { + name: "observed count wins over the estimator", + events: []*session.Event{ + withUsage(modelTextEvent("a", "inv1", 1, "a1"), 500), + }, + estimate: func([]*session.Event) int { return 123 }, + want: 500, + wantOK: true, + }, + { + name: "the most recent observed count wins", + events: []*session.Event{ + withUsage(modelTextEvent("a", "inv1", 1, "a1"), 500), + textEvent("b", "inv2", 2, "q2"), + withUsage(modelTextEvent("c", "inv2", 3, "a2"), 900), + }, + want: 900, + wantOK: true, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + got, ok := promptTokenCount(tc.events, TurnScope{}, tc.estimate) + if got != tc.want || ok != tc.wantOK { + t.Errorf("promptTokenCount() = (%d, TurnScope{}, %t), want (%d, %t)", got, ok, tc.want, tc.wantOK) + } + }) + } +} + +func TestEstimateTokensFromContents(t *testing.T) { + t.Parallel() + + text := func(n int) *genai.Content { + return &genai.Content{Parts: []*genai.Part{{Text: strings.Repeat("x", n)}}} + } + + tests := []struct { + name string + contents []*genai.Content + want int + }{ + {name: "nil", contents: nil, want: 0}, + {name: "empty text", contents: []*genai.Content{text(0)}, want: 0}, + {name: "below one token", contents: []*genai.Content{text(3)}, want: 0}, + {name: "exactly one token", contents: []*genai.Content{text(4)}, want: 1}, + {name: "summed across contents", contents: []*genai.Content{text(2000), text(2000)}, want: 1000}, + {name: "nil content is skipped", contents: []*genai.Content{nil, text(4)}, want: 1}, + {name: "nil part is skipped", contents: []*genai.Content{{Parts: []*genai.Part{nil, {Text: "xxxx"}}}}, want: 1}, + { + // Tool traffic counts. It is what a long turn grows by, and the + // estimate exists to notice a long turn growing: "search" is six + // characters, so it is a token and a half's worth on its own. + name: "a function call counts its name", + contents: []*genai.Content{{Parts: []*genai.Part{{FunctionCall: &genai.FunctionCall{Name: "search"}}}}}, + want: 1, + }, + { + // The payload dominates, and counting only Text saw none of it. + name: "a function response counts its payload", + contents: []*genai.Content{{Parts: []*genai.Part{{ + FunctionResponse: &genai.FunctionResponse{ + Name: "search", + Response: map[string]any{"result": strings.Repeat("y", 4000)}, + }, + }}}}, + // 4000 characters of payload, so a thousand tokens give or take the + // JSON punctuation and the name. + want: 1004, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + if got := EstimateTokensFromContents(tc.contents); got != tc.want { + t.Errorf("EstimateTokensFromContents() = %d, want %d", got, tc.want) + } + }) + } +} + +func TestTailRetention(t *testing.T) { + t.Parallel() + + fourEvents := []*session.Event{ + textEvent("a", "inv1", 1, "q1"), modelTextEvent("b", "inv1", 2, "a1"), + textEvent("c", "inv2", 3, "q2"), withUsage(modelTextEvent("d", "inv2", 4, "a2"), 900), + } + + tests := []struct { + name string + cfg *compaction.Config + events []*session.Event + summarizer *fakeSummarizer + wantSummary bool + wantWindow []string + wantErr bool + }{ + { + name: "nil config does nothing", + cfg: nil, + events: fourEvents, + summarizer: &fakeSummarizer{summary: "sum"}, + }, + { + name: "sliding-window-only config does nothing", + cfg: &compaction.Config{CompactionInterval: 2}, + events: fourEvents, + summarizer: &fakeSummarizer{summary: "sum"}, + }, + { + name: "below the threshold", + cfg: &compaction.Config{TokenThreshold: 1000, EventRetentionSize: 2}, + events: fourEvents, + summarizer: &fakeSummarizer{summary: "sum"}, + }, + { + name: "at the threshold", + cfg: &compaction.Config{TokenThreshold: 900, EventRetentionSize: 2}, + events: fourEvents, + summarizer: &fakeSummarizer{summary: "sum"}, + wantSummary: true, + wantWindow: []string{"a", "b"}, + }, + { + name: "above the threshold", + cfg: &compaction.Config{TokenThreshold: 100, EventRetentionSize: 2}, + events: fourEvents, + summarizer: &fakeSummarizer{summary: "sum"}, + wantSummary: true, + wantWindow: []string{"a", "b"}, + }, + { + name: "threshold reached but the tail retains everything", + cfg: &compaction.Config{TokenThreshold: 100, EventRetentionSize: 10}, + events: fourEvents, + summarizer: &fakeSummarizer{summary: "sum"}, + }, + { + name: "summarizer declines", + cfg: &compaction.Config{TokenThreshold: 100, EventRetentionSize: 2}, + events: fourEvents, + summarizer: &fakeSummarizer{}, + wantSummary: false, + wantWindow: []string{"a", "b"}, + }, + { + name: "summarizer fails", + cfg: &compaction.Config{TokenThreshold: 100, EventRetentionSize: 2}, + events: fourEvents, + summarizer: &fakeSummarizer{err: errors.New("boom")}, + wantWindow: []string{"a", "b"}, + wantErr: true, + }, + { + name: "no observed token count and no estimate", + cfg: &compaction.Config{TokenThreshold: 1, EventRetentionSize: 1}, + events: []*session.Event{textEvent("a", "inv1", 1, "q1"), textEvent("b", "inv1", 2, "q2")}, + summarizer: &fakeSummarizer{summary: "sum"}, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + cfg := tc.cfg + if cfg != nil { + copied := *cfg + copied.Summarizer = tc.summarizer + cfg = &copied + } + + got, err := tailRetentionStored(context.Background(), cfg, &staticSession{events: tc.events}, TurnScope{}, nil, nil) + if gotErr := err != nil; gotErr != tc.wantErr { + t.Fatalf("tailRetentionStored() error = %v, wantErr %t", err, tc.wantErr) + } + if gotSummary := got != nil; gotSummary != tc.wantSummary { + t.Errorf("tailRetentionStored() returned event = %t, want %t", gotSummary, tc.wantSummary) + } + var gotWindow []string + if len(tc.summarizer.windows) > 0 { + gotWindow = tc.summarizer.windows[0] + } + if diff := cmp.Diff(tc.wantWindow, gotWindow); diff != "" { + t.Errorf("summarizer window mismatch (-want +got):\n%s", diff) + } + }) + } +} + +func TestTailRetentionUsesTheEstimator(t *testing.T) { + t.Parallel() + + // No event carries usage metadata, so the estimator decides. + events := []*session.Event{ + textEvent("a", "inv1", 1, "q1"), modelTextEvent("b", "inv1", 2, "a1"), + textEvent("c", "inv2", 3, "q2"), modelTextEvent("d", "inv2", 4, "a2"), + } + summarizer := &fakeSummarizer{summary: "sum"} + cfg := &compaction.Config{TokenThreshold: 500, EventRetentionSize: 2, Summarizer: summarizer} + + got, err := tailRetentionStored(context.Background(), cfg, &staticSession{events: events}, TurnScope{}, + func([]*session.Event) int { return 100 }, nil) + if err != nil { + t.Fatalf("tailRetentionStored() error = %v", err) + } + if got != nil { + t.Error("tailRetentionStored() compacted despite an estimate below the threshold") + } + + got, err = tailRetentionStored(context.Background(), cfg, &staticSession{events: events}, TurnScope{}, + func([]*session.Event) int { return 700 }, nil) + if err != nil { + t.Fatalf("tailRetentionStored() error = %v", err) + } + if got == nil { + t.Error("tailRetentionStored() did not compact despite an estimate above the threshold") + } +} + +func TestTailRetentionRequiresSummarizer(t *testing.T) { + t.Parallel() + + _, err := tailRetentionStored(context.Background(), &compaction.Config{TokenThreshold: 1, EventRetentionSize: 0}, + &staticSession{events: []*session.Event{withUsage(modelTextEvent("a", "inv1", 1, "a"), 10)}}, TurnScope{}, nil, nil) + if err == nil { + t.Fatal("tailRetentionStored() with no Summarizer returned nil error, want an error") + } +} + +func TestTailRetentionStampsTheSummary(t *testing.T) { + t.Parallel() + + events := []*session.Event{ + textEvent("a", "inv1", 1, "q1"), + withUsage(modelTextEvent("b", "inv1", 2, "a1"), 900), + } + cfg := &compaction.Config{TokenThreshold: 100, EventRetentionSize: 0, Summarizer: &fakeSummarizer{summary: "sum"}} + + got, err := tailRetentionStored(context.Background(), cfg, &staticSession{events: events}, TurnScope{}, nil, nil) + if err != nil { + t.Fatalf("tailRetentionStored() error = %v", err) + } + if got == nil { + t.Fatal("tailRetentionStored() produced no summary") + } + // The event must be ready to append without the caller filling anything in. + if got.ID == "" { + t.Error("summary has no ID") + } + if got.InvocationID == "" { + t.Error("summary has no InvocationID") + } + if got.Timestamp.IsZero() { + t.Error("summary has no Timestamp") + } + for _, ev := range events { + if got.InvocationID == ev.InvocationID { + t.Errorf("summary reuses invocation ID %q from a covered event; window selection counts invocations, so it must be fresh", got.InvocationID) + } + } +} + +// TestTailRetentionThenApplyShrinksHistory is the round trip: compact, then +// build the prompt, and confirm the covered events are gone. +func TestTailRetentionThenApplyShrinksHistory(t *testing.T) { + t.Parallel() + + events := []*session.Event{ + textEvent("a", "inv1", 1, "q1"), modelTextEvent("b", "inv1", 2, "a1"), + textEvent("c", "inv1", 3, "q2"), withUsage(modelTextEvent("d", "inv1", 4, "a2"), 5000), + } + cfg := &compaction.Config{TokenThreshold: 1000, EventRetentionSize: 2, Summarizer: &fakeSummarizer{summary: "SUMMARY"}} + + summary, err := tailRetentionStored(context.Background(), cfg, &staticSession{events: events}, TurnScope{}, nil, nil) + if err != nil { + t.Fatalf("tailRetentionStored() error = %v", err) + } + if summary == nil { + t.Fatal("tailRetentionStored() produced no summary") + } + summary.ID = "s1" + + got := Apply(append(events, summary)) + if diff := cmp.Diff([]string{"s1", "c", "d"}, ids(got)); diff != "" { + t.Errorf("post-compaction prompt events mismatch (-want +got):\n%s", diff) + } + if texts := utils.TextParts(utils.Content(got[0])); len(texts) != 1 || texts[0] != "SUMMARY" { + t.Errorf("first prompt event = %v, want the summary text", texts) + } +} + +// TestSelectTailRetentionWindowStaysInOneScope checks that the tail window stops +// at the first branch or isolation-scope change. +// +// A summary inherits the branch and isolation scope of what it covers, so a +// window spanning two of them produces one summary that necessarily misattributes +// half its content. Stamped with the first event's scope, it becomes readable by +// agents the filters exist to keep the rest away from. +func TestSelectTailRetentionWindowStaysInOneScope(t *testing.T) { + t.Parallel() + + root1 := textEvent("a", "inv1", 1, "q1") + root2 := modelTextEvent("b", "inv1", 2, "a1") + sub := textEvent("c", "inv2", 3, "SUB-AGENT-SECRET") + sub.Branch = "root.sub" + sub.IsolationScope = "scope-1" + tail1 := textEvent("d", "inv3", 4, "q3") + tail2 := modelTextEvent("e", "inv3", 5, "a3") + + events := []*session.Event{root1, root2, sub, tail1, tail2} + + window := selectTailRetentionWindow(events, 2, TurnScope{}) + if diff := cmp.Diff([]string{"a", "b"}, ids(window)); diff != "" { + t.Errorf("selectTailRetentionWindow() mismatch (-want +got):\n%s\nthe window must stop at the scope change", diff) + } + for _, ev := range window { + if ev.Branch != "" || ev.IsolationScope != "" { + t.Errorf("event %q carries branch %q scope %q, so the window is not homogeneous", ev.ID, ev.Branch, ev.IsolationScope) + } + } +} + +// TestSelectTailRetentionWindowKeepsATiedBoundaryEvent checks that an event +// stamped exactly at the previous compaction's end is not lost. +// +// The candidate filter used to exclude anything not strictly after that +// instant, while the new range, seeded with the previous summary, starts back +// at the previous start and so covers it. An event on that boundary therefore +// went into no window and inside the next recorded range: summarized by +// nothing, and dropped from every prompt afterwards. +func TestSelectTailRetentionWindowKeepsATiedBoundaryEvent(t *testing.T) { + t.Parallel() + + prior := compactionEvent("s1", 3, 1, 3, "EARLIER") + // Appended after the compaction, but stamped on its end instant. + tied := textEvent("tied", "inv2", 3, "NEVER-SUMMARIZED") + events := []*session.Event{ + textEvent("a", "inv1", 1, "q1"), + modelTextEvent("b", "inv1", 2, "a1"), + prior, + tied, + textEvent("c", "inv3", 4, "q3"), + modelTextEvent("d", "inv3", 5, "a3"), + textEvent("e", "inv4", 6, "q4"), + } + + window := selectTailRetentionWindow(events, 1, TurnScope{}) + if !slices.Contains(ids(window), "tied") { + t.Errorf("window %v does not include the boundary event, so it is covered by the next range without being summarized", ids(window)) + } +} + +// TestPromptTokenCountAddsEventsSinceTheLastReport checks that the count is not +// stale by a whole turn. +// +// A reported count describes the prompt of an earlier call. Returning it +// unchanged means everything appended since is invisible, so the call that +// first crosses the threshold is missed and compaction reacts one call late. +func TestPromptTokenCountAddsEventsSinceTheLastReport(t *testing.T) { + t.Parallel() + + reported := modelTextEvent("a", "inv1", 1, "answer") + reported.LLMResponse.UsageMetadata = &genai.GenerateContentResponseUsageMetadata{PromptTokenCount: 100} + events := []*session.Event{ + reported, + textEvent("b", "inv2", 2, strings.Repeat("x", 400)), + } + + // The estimator stands in for the real one: four characters per token. + estimate := func(evs []*session.Event) int { + n := 0 + for _, ev := range evs { + for _, p := range utils.Content(ev).Parts { + n += len(p.Text) + } + } + return n / 4 + } + + got, ok := promptTokenCount(events, TurnScope{}, estimate) + if !ok { + t.Fatal("promptTokenCount() reported nothing") + } + if got <= 100 { + t.Errorf("promptTokenCount() = %d, TurnScope{}, want more than the reported 100: the 400 characters appended since are not counted", got) + } +} + +// recordingGate captures the [ProgressGate] calls TailRetention makes. +type recordingGate struct { + allow bool + recorded []int + recovered int + failed int +} + +func (g *recordingGate) AllowAt() bool { return g.allow } +func (g *recordingGate) RecordAt(t int) { g.recorded = append(g.recorded, t) } +func (g *recordingGate) Recovered() { g.recovered++ } +func (g *recordingGate) Failed() { g.failed++ } + +// TestTailRetentionReArmsTheGateBelowTheThreshold pins that a prompt back under +// the threshold re-arms the gate. +// +// Without this the gate closes on the first compaction of a turn and never +// reopens, so a long turn that keeps growing never compacts again. +func TestTailRetentionReArmsTheGateBelowTheThreshold(t *testing.T) { + t.Parallel() + + events := []*session.Event{ + textEvent("a", "inv1", 1, "q1"), modelTextEvent("b", "inv1", 2, "a1"), + textEvent("c", "inv2", 3, "q2"), withUsage(modelTextEvent("d", "inv2", 4, "a2"), 100), + } + gate := &recordingGate{allow: true} + cfg := &compaction.Config{TokenThreshold: 1000, EventRetentionSize: 2, Summarizer: &fakeSummarizer{summary: "sum"}} + + got, err := tailRetentionStored(context.Background(), cfg, &staticSession{events: events}, TurnScope{}, nil, gate) + if err != nil { + t.Fatalf("tailRetentionStored() error = %v", err) + } + if got != nil { + t.Fatalf("tailRetentionStored() returned a summary at 100 tokens against a 1000 threshold") + } + if gate.recovered != 1 { + t.Errorf("Recovered() called %d times, want 1: a prompt under the threshold means the last compaction worked", gate.recovered) + } +} + +// TestTailRetentionDoesNotRecordAFailedAttempt pins that a summarizer failure +// leaves the progress gate as it found it. +// +// Recording the attempt rather than the result let one transient error disarm +// compaction for the whole invocation with nothing stored in exchange, and the +// prompt then grew unchecked behind a gate that had stopped retrying. +func TestTailRetentionDoesNotRecordAFailedAttempt(t *testing.T) { + t.Parallel() + + events := []*session.Event{ + textEvent("a", "inv1", 1, "q1"), modelTextEvent("b", "inv1", 2, "a1"), + textEvent("c", "inv2", 3, "q2"), withUsage(modelTextEvent("d", "inv2", 4, "a2"), 900), + } + gate := &recordingGate{allow: true} + cfg := &compaction.Config{TokenThreshold: 100, EventRetentionSize: 2, Summarizer: &fakeSummarizer{err: errors.New("boom")}} + + if _, err := tailRetentionStored(context.Background(), cfg, &staticSession{events: events}, TurnScope{}, nil, gate); err == nil { + t.Fatal("tailRetentionStored() error = nil, want the summarizer failure") + } + if len(gate.recorded) != 0 { + t.Errorf("RecordAt called %v after a failed summarization, want no calls", gate.recorded) + } +} + +// TestTailRetentionRecordsASuccessfulCompaction is the counterpart: a summary +// that was produced must close the gate. +func TestTailRetentionRecordsASuccessfulCompaction(t *testing.T) { + t.Parallel() + + events := []*session.Event{ + textEvent("a", "inv1", 1, "q1"), modelTextEvent("b", "inv1", 2, "a1"), + textEvent("c", "inv2", 3, "q2"), withUsage(modelTextEvent("d", "inv2", 4, "a2"), 900), + } + gate := &recordingGate{allow: true} + cfg := &compaction.Config{TokenThreshold: 100, EventRetentionSize: 2, Summarizer: &fakeSummarizer{summary: "sum"}} + + got, err := tailRetentionStored(context.Background(), cfg, &staticSession{events: events}, TurnScope{}, nil, gate) + if err != nil || got == nil { + t.Fatalf("tailRetentionStored() = %v, %v, want a summary and no error", got, err) + } + if diff := cmp.Diff([]int{900}, gate.recorded); diff != "" { + t.Errorf("RecordAt calls mismatch (-want +got):\n%s", diff) + } +} + +// TestSelectTailRetentionWindowKeepsTheLiveQuestion pins that the turn being +// answered keeps its own question. +// +// EventRetentionSize counts events and a turn is not a fixed number of them, so +// at every size Validate accepts the question can scroll out of the retained +// tail and be summarized into a paraphrase of the instruction being carried +// out. It is held back separately. +// +// The traffic after it stays eligible, which is the point: excluding the whole +// live invocation would stop a long tool loop compacting itself, and that is +// the case this strategy exists for. +func TestSelectTailRetentionWindowKeepsTheLiveQuestion(t *testing.T) { + t.Parallel() + + events := []*session.Event{ + textEvent("q1", "inv1", 1, "older question"), + modelTextEvent("a1", "inv1", 2, "older answer"), + // The turn in flight: its question, then a long tool loop. + textEvent("q2", "inv2", 3, "the question being answered"), + modelTextEvent("t1", "inv2", 4, "tool step 1"), + modelTextEvent("t2", "inv2", 5, "tool step 2"), + modelTextEvent("t3", "inv2", 6, "tool step 3"), + } + + got := ids(selectTailRetentionWindow(events, 2, TurnScope{InvocationID: "inv2"})) + + if slices.Contains(got, "q2") { + t.Error("the window covers the question the turn is answering") + } + // The loop's own older traffic is still compactable, skipping over the + // question, which only a covered set can express. + if diff := cmp.Diff([]string{"q1", "a1", "t1"}, got); diff != "" { + t.Errorf("window mismatch (-want +got):\n%s", diff) + } +} + +// TestSelectTailRetentionWindowStepsPastABlockedHead pins that one unanswered +// tool call does not stop tail retention for the rest of the session. +// +// The window is anchored to the last compaction boundary, so a call awaiting +// human approval, or one whose backend died, sits at the head of every later +// attempt. The sliding window already steps past it; tail retention gave up +// instead, and gave up silently, since "no self-contained prefix" and "nothing +// to do" both come back as nil. Long tool-using sessions are exactly the ones +// this strategy exists for. +func TestSelectTailRetentionWindowStepsPastABlockedHead(t *testing.T) { + t.Parallel() + + events := []*session.Event{ + // A call at the head that nothing ever answers. + callEvent("blocked", "inv1", 1, "c-pending"), + // A complete exchange behind it, which is compactable. + callEvent("call", "inv2", 2, "c-done"), + responseEvent("resp", "inv2", 3, "c-done"), + textEvent("q", "inv3", 4, "another question"), + modelTextEvent("a", "inv3", 5, "another answer"), + } + + got := ids(selectTailRetentionWindow(events, 2, TurnScope{})) + if len(got) == 0 { + t.Fatal("selectTailRetentionWindow() gave up because the head is blocked") + } + if slices.Contains(got, "blocked") { + t.Error("the window covers the pending call, which must stay raw and visible") + } + if diff := cmp.Diff([]string{"call", "resp"}, got); diff != "" { + t.Errorf("window mismatch (-want +got):\n%s", diff) + } +} + +// TestSkipBlockedHeadKeepsACallWithItsResponse pins that stepping past a +// blocked head never summarizes a response whose call stays raw. +// +// longestSelfContainedPrefix only tracks obligations opened inside the slice it +// is handed, so a response whose call sits in the skipped head looked +// unremarkable: the response was summarized, the call stayed raw, and the model +// was shown a call it had already answered with the answer gone. +func TestSkipBlockedHeadKeepsACallWithItsResponse(t *testing.T) { + t.Parallel() + + window := []*session.Event{ + // One event opening two calls: the head is blocked on c-pending, and + // c-two is answered below. Any resume point is therefore past both + // calls, so the response is the first thing the tail sees. + multiCallEvent("head", "inv1", 1, "c-pending", "c-two"), + responseEvent("resp2", "inv1", 2, "c-two"), + textEvent("q", "inv2", 3, "later question"), + modelTextEvent("a", "inv2", 4, "later answer"), + } + + got := ids(skipBlockedHead(window)) + if slices.Contains(got, "resp2") { + t.Errorf("window %v summarizes a response whose call stays raw in the skipped head", got) + } +} + +// TestPromptTokenCountIgnoresOtherBranches pins that a turn reads a token count +// describing its own prompt. +// +// The count decides whether this turn's prompt is too large, and that prompt is +// assembled with branch and isolation-scope filtering. Reading the newest count +// from anywhere in the session meant a sub-agent whose own prompt is a few +// tokens inherited its parent's, and compacted history it had no business +// compacting. +func TestPromptTokenCountIgnoresOtherBranches(t *testing.T) { + t.Parallel() + + onBranch := func(ev *session.Event, branch string) *session.Event { + ev.Branch = branch + return ev + } + events := []*session.Event{ + withUsage(onBranch(modelTextEvent("mine", "inv1", 1, "small"), "parent.child"), 40), + // A sibling's turn, invisible to parent.child, reporting a huge prompt. + withUsage(onBranch(modelTextEvent("sibling", "inv2", 2, "huge"), "parent.other"), 200000), + } + + got, ok := promptTokenCount(events, TurnScope{Branch: "parent.child"}, nil) + if !ok { + t.Fatal("promptTokenCount() found no count at all") + } + if got != 40 { + t.Errorf("promptTokenCount() = %d, want 40: the reading came from another branch", got) + } +} + +// TestPromptTokenCountIgnoresOtherIsolationScopes is the same property for +// isolation scope, which is an exact match rather than an ancestor one. +func TestPromptTokenCountIgnoresOtherIsolationScopes(t *testing.T) { + t.Parallel() + + scoped := func(ev *session.Event, scope string) *session.Event { + ev.IsolationScope = scope + return ev + } + events := []*session.Event{ + withUsage(scoped(modelTextEvent("mine", "inv1", 1, "small"), "task-a"), 40), + withUsage(scoped(modelTextEvent("other", "inv2", 2, "huge"), "task-b"), 200000), + } + + got, ok := promptTokenCount(events, TurnScope{IsolationScope: "task-a"}, nil) + if !ok { + t.Fatal("promptTokenCount() found no count at all") + } + if got != 40 { + t.Errorf("promptTokenCount() = %d, want 40: the reading came from another isolation scope", got) + } +} + +// TestTailRetentionDoesNotRecordADiscardedSummary pins that a summary the +// caller throws away leaves the progress gate as it found it. +// +// Recording the attempt rather than the result disarmed compaction for the rest +// of an invocation with nothing stored. Moving the call past the summarizer +// fixed the transient-error case and left four others, because the caller can +// still discard a perfectly good summary: a cancelled turn, a failed re-read, a +// competing compaction, a failed append. Each closed the gate on a summary that +// never existed, and Recovered cannot reopen it because the prompt never drops. +func TestTailRetentionDoesNotRecordADiscardedSummary(t *testing.T) { + t.Parallel() + + events := []*session.Event{ + textEvent("a", "inv1", 1, "q1"), modelTextEvent("b", "inv1", 2, "a1"), + textEvent("c", "inv2", 3, "q2"), withUsage(modelTextEvent("d", "inv2", 4, "a2"), 900), + } + tests := []struct { + name string + finishErr error + discardReason string + wantRecorded bool + }{ + {name: "stored", wantRecorded: true}, + {name: "discarded by the caller", discardReason: "a competing compaction landed"}, + {name: "failed on the way to the session", finishErr: errors.New("append failed")}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + // Per subtest: the summarizer counts its calls, so sharing one + // across parallel subtests races. + cfg := &compaction.Config{ + TokenThreshold: 100, EventRetentionSize: 2, + Summarizer: &fakeSummarizer{summary: "sum"}, + } + gate := &recordingGate{allow: true} + summary, finish, err := TailRetention(context.Background(), cfg, + &staticSession{events: events}, TurnScope{}, nil, gate) + if err != nil || summary == nil { + t.Fatalf("TailRetention() = %v, %v, want a summary", summary, err) + } + finish(tt.finishErr, tt.discardReason) + + if gotRecorded := len(gate.recorded) > 0; gotRecorded != tt.wantRecorded { + t.Errorf("gate recorded = %t, want %t: %v", gotRecorded, tt.wantRecorded, gate.recorded) + } + }) + } +} + +// TestSkipBlockedHeadStillCompactsPastAnAnsweredSibling pins the positive half +// of TestSkipBlockedHeadKeepsACallWithItsResponse, which only asserts that a +// response is not summarized and is therefore satisfied by giving up entirely. +// +// One model turn emitting two calls, one to an ordinary tool and one to a +// long-running tool that never produces a response, is the standard +// long-running shape. The answered sibling's response necessarily sits after +// both calls, so every resume point the scan was willing to consider had that +// response in the tail and a call for it still open in the head, and all of +// them were refused. Nothing after the blockage was ever compacted again, for +// the rest of the session, and because "no window" and "nothing to do yet" are +// both nil it was silent. +// +// The resume point that works is the one just after the response: the head then +// holds the call and its answer, only the long-running call is still open, and +// the tail answers nothing. The scan skipped it because it only resumed after +// an event that opened an obligation. +func TestSkipBlockedHeadStillCompactsPastAnAnsweredSibling(t *testing.T) { + t.Parallel() + + window := []*session.Event{ + multiCallEvent("head", "inv1", 1, "c-longrunning", "c-two"), + responseEvent("resp2", "inv1", 2, "c-two"), + textEvent("q", "inv2", 3, "later question"), + modelTextEvent("a", "inv2", 4, "later answer"), + textEvent("q2", "inv3", 5, "later question 2"), + modelTextEvent("a2", "inv3", 6, "later answer 2"), + } + + got := ids(skipBlockedHead(window)) + if len(got) == 0 { + t.Fatal("skipBlockedHead() = nil: one unanswered call in a parallel pair stalls compaction for the rest of the session") + } + if slices.Contains(got, "resp2") { + t.Errorf("window %v summarizes a response whose call stays raw in the skipped head", got) + } + if diff := cmp.Diff([]string{"q", "a", "q2", "a2"}, got); diff != "" { + t.Errorf("window mismatch (-want +got):\n%s", diff) + } +} + +// TestTailRetentionClosesTheGateOnlyWhenTryingAgainIsPointless pins which +// outcomes stop this invocation attempting again, and which do not. +// +// The gate has to sit between two failure modes that pull opposite ways. +// Closing it on every attempt disarms compaction on a transient error and +// Recovered cannot reopen it, because the prompt never drops. Closing it only +// on a stored summary means a persistently failing summarizer is retried before +// every model call for the life of the turn. +// +// A discard is neither. The commonest one is a competing compaction landing in +// the range, which means a summary was stored and the picture changed, so the +// next call should look again. +func TestTailRetentionClosesTheGateOnlyWhenTryingAgainIsPointless(t *testing.T) { + t.Parallel() + + events := []*session.Event{ + withUsage(modelTextEvent("a", "inv1", 1, "a1"), 900), + textEvent("b", "inv1", 2, "q2"), + modelTextEvent("c", "inv1", 3, "a2"), + textEvent("d", "inv1", 4, "q3"), + } + run := func(t *testing.T, sum compaction.Summarizer, finishWith func(finish Finish)) *recordingGate { + t.Helper() + gate := &recordingGate{allow: true} + cfg := &compaction.Config{TokenThreshold: 100, EventRetentionSize: 1, Summarizer: sum} + _, finish, err := TailRetention(context.Background(), cfg, &staticSession{events: events}, + TurnScope{}, nil, gate) + if err == nil && finishWith != nil { + finishWith(finish) + } + return gate + } + + t.Run("the summarizer errored", func(t *testing.T) { + t.Parallel() + // The caller never gets a Finish on this path, so closing the gate only + // in that callback did not reach here at all. + g := run(t, &fakeSummarizer{err: errTestSummarizer}, nil) + if g.failed != 1 { + t.Errorf("gate.failed = %d, want 1: a failing summarizer must not be retried every model call", g.failed) + } + }) + t.Run("the summarizer declined", func(t *testing.T) { + t.Parallel() + g := run(t, &fakeSummarizer{}, nil) + if g.failed != 1 { + t.Errorf("gate.failed = %d, want 1: a decline gets no Finish either", g.failed) + } + }) + t.Run("the append failed", func(t *testing.T) { + t.Parallel() + g := run(t, &fakeSummarizer{summary: "a summary"}, func(f Finish) { f(errTestAppend, "") }) + if g.failed != 1 { + t.Errorf("gate.failed = %d, want 1", g.failed) + } + if len(g.recorded) != 0 { + t.Errorf("gate recorded %v for a summary that was never stored", g.recorded) + } + }) + t.Run("the caller discarded it", func(t *testing.T) { + t.Parallel() + g := run(t, &fakeSummarizer{summary: "a summary"}, func(f Finish) { f(nil, "a competing compaction landed") }) + if g.failed != 0 { + t.Errorf("gate.failed = %d, want 0: a discard means the picture changed, not that trying again is pointless", g.failed) + } + if len(g.recorded) != 0 { + t.Errorf("gate recorded %v for a summary that was never stored", g.recorded) + } + }) + t.Run("it was stored", func(t *testing.T) { + t.Parallel() + g := run(t, &fakeSummarizer{summary: "a summary"}, func(f Finish) { f(nil, "") }) + if len(g.recorded) != 1 { + t.Errorf("gate.recorded = %v, want one entry", g.recorded) + } + if g.failed != 0 { + t.Errorf("gate.failed = %d, want 0", g.failed) + } + }) +} + +var ( + errTestAppend = errors.New("append failed") + errTestSummarizer = errors.New("summarizer unavailable") +) + +// TestSelectTailRetentionWindowSeedsWhileHoldingBackTheLiveHead exercises the +// two mechanisms together, which is the combination production always runs and +// the suite never covered. +// +// The seed keeps history to one rolling summary. The live-head hold-back keeps +// the current turn's question out of the window. Every other case in this file +// passes an empty TurnScope, so the hold-back is inert in all of them, while +// the processor always supplies a real InvocationID. The two were each +// exercised alone and never together, which is how a defect in their +// interaction could sit here unnoticed. +func TestSelectTailRetentionWindowSeedsWhileHoldingBackTheLiveHead(t *testing.T) { + t.Parallel() + + events := []*session.Event{ + textEvent("a", "inv1", 1, "q1"), modelTextEvent("b", "inv1", 2, "a1"), + compactionEvent("s1", 3, 1, 2, "earlier summary"), + textEvent("c", "inv2", 4, "q2"), modelTextEvent("d", "inv2", 5, "a2"), + // The live turn: its question must stay out, and the events after it + // are fair game. + textEvent("live-q", "inv3", 6, "the question being answered now"), + modelTextEvent("live-a", "inv3", 7, "partial work"), + textEvent("live-b", "inv3", 8, "more partial work"), + } + + window := selectTailRetentionWindow(events, 1, TurnScope{InvocationID: "inv3"}) + got := ids(window) + + if slices.Contains(got, "live-q") { + t.Errorf("window %v summarizes the question the turn is answering", got) + } + if len(got) == 0 || got[0] != "s1" { + t.Errorf("window %v does not open with the previous summary, so it cannot supersede it", got) + } + + // And the record built from it supersedes the one it was seeded with. + summary, err := newSummaryEvent(window, events, genai.NewContentFromText("new summary", "model"), nil) + if err != nil { + t.Fatalf("newSummaryEvent() error = %v", err) + } + summary.ID, summary.Timestamp = "s2", at(9) + after := ids(Apply(append(events, summary))) + if slices.Contains(after, "s1") { + t.Errorf("prompt %v still carries the superseded summary alongside its replacement", after) + } + if !slices.Contains(after, "live-q") { + t.Errorf("prompt %v lost the question the turn is answering", after) + } +} + +// TestTailRetentionEmptyWindowIsNotYetNotNever pins that finding nothing to +// summarize does not stop the invocation trying later. +// +// The window is empty when the retained tail is still the whole history, or +// when a tool call at its head is unanswered. A long tool-calling turn keeps +// appending, so a later model call in the same invocation can have a window +// when this one does not. Treating the empty window as a spent attempt stopped +// compaction for the whole turn on the first cheap check, before the summarizer +// had been asked anything. +// +// The distinction that matters: this path makes no model call, so re-checking +// is free, where retrying a failing summarizer is not. +func TestTailRetentionEmptyWindowIsNotYetNotNever(t *testing.T) { + t.Parallel() + + gate := &recordingGate{allow: true} + // One event over the threshold and a retention size that keeps it, so + // there is nothing left to summarize. + events := []*session.Event{withUsage(modelTextEvent("a", "inv1", 1, "a1"), 900)} + cfg := &compaction.Config{TokenThreshold: 100, EventRetentionSize: 4, Summarizer: &fakeSummarizer{summary: "s"}} + + summary, _, err := TailRetention(context.Background(), cfg, &staticSession{events: events}, + TurnScope{}, nil, gate) + if err != nil || summary != nil { + t.Fatalf("TailRetention() = %v, %v, want no summary and no error", summary, err) + } + if gate.failed != 0 { + t.Errorf("gate.failed = %d, want 0: an empty window is not a spent attempt", gate.failed) + } +} + +// TestSeedIsTakenFromTheWindowsOwnBranch pins that a sibling branch compacting +// more recently does not stop a branch rolling up its own summary. +// +// The seed keeps history to one rolling summary. Taking the latest record +// across every branch and then refusing it on a scope mismatch looks safe and +// is not the same thing: when a sibling compacted last, the seed is dropped and +// this branch never absorbs its own earlier summary, so it accumulates a chain +// and the bound quietly stops applying. Nothing is lost, which is why it went +// unnoticed. +func TestSeedIsTakenFromTheWindowsOwnBranch(t *testing.T) { + t.Parallel() + + onBranch := func(ev *session.Event, branch string) *session.Event { + ev.Branch = branch + return ev + } + events := []*session.Event{ + onBranch(textEvent("a1", "invA1", 1, "A one"), "root.A"), + onBranch(modelTextEvent("a2", "invA1", 2, "A two"), "root.A"), + onBranch(compactionEvent("SA", 3, 1, 2, "summary of A"), "root.A"), + onBranch(textEvent("b1", "invB1", 4, "B one"), "root.B"), + onBranch(compactionEvent("SB", 5, 4, 4, "summary of B"), "root.B"), + onBranch(textEvent("a3", "invA2", 6, "A three"), "root.A"), + onBranch(modelTextEvent("a4", "invA2", 7, "A four"), "root.A"), + onBranch(textEvent("a5", "invA3", 8, "A five"), "root.A"), + } + + window := selectTailRetentionWindow(events, 1, TurnScope{Branch: "root.A"}) + got := ids(window) + if !slices.Contains(got, "SA") { + t.Errorf("window %v does not open with branch A's own summary, so A never rolls it up and grows a chain", got) + } + if slices.Contains(got, "SB") { + t.Errorf("window %v seeded from a sibling branch, merging across the scope boundary", got) + } +} + +// TestTailRetentionWindowStaysInsideTheTurnsScope pins that the window is drawn +// from the same conversation whose size triggered the pass. +// +// promptTokenCount measures only what the turn can see, and its comment gives +// the reason: "a sub-agent whose own prompt is a couple of tokens read its +// parent's 200,000 and compacted history it had no business compacting". The +// candidate list did not apply that filter, so the threshold was evaluated +// against one conversation and the window chosen from another. A sub-agent +// spent its one allowed compaction summarizing a sibling branch, its own prompt +// did not shrink by a token, and the progress gate then stood down for the rest +// of an over-threshold turn -- the strategy that bounds growth, switched off by +// having compacted the wrong thing. +func TestTailRetentionWindowStaysInsideTheTurnsScope(t *testing.T) { + t.Parallel() + + onBranch := func(ev *session.Event, branch string) *session.Event { + ev.Branch = branch + return ev + } + events := []*session.Event{ + onBranch(textEvent("b1", "invB", 1, "the sibling's only event"), "root.B"), + onBranch(textEvent("a1", "invA1", 2, "A one"), "root.A"), + onBranch(modelTextEvent("a2", "invA1", 3, "A two"), "root.A"), + onBranch(textEvent("a3", "invA2", 4, "A three"), "root.A"), + onBranch(modelTextEvent("a4", "invA2", 5, "A four"), "root.A"), + } + + got := ids(selectTailRetentionWindow(events, 2, TurnScope{Branch: "root.A"})) + + if len(got) == 0 { + t.Fatal("selectTailRetentionWindow() = empty, want root.A's own eligible events") + } + if slices.Contains(got, "b1") { + t.Errorf("window %v summarizes a sibling branch: root.A's prompt does not shrink, and the compaction it was allowed is spent", got) + } + if diff := cmp.Diff([]string{"a1", "a2"}, got); diff != "" { + t.Errorf("window (-want +got):\n%s", diff) + } +} diff --git a/internal/compactioninternal/telemetry_test.go b/internal/compactioninternal/telemetry_test.go new file mode 100644 index 000000000..f9161b1f1 --- /dev/null +++ b/internal/compactioninternal/telemetry_test.go @@ -0,0 +1,813 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package compactioninternal + +import ( + "context" + "errors" + "slices" + "testing" + "time" + + "github.com/google/go-cmp/cmp" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/codes" + sdktrace "go.opentelemetry.io/otel/sdk/trace" + "go.opentelemetry.io/otel/sdk/trace/tracetest" + "google.golang.org/genai" + + "google.golang.org/adk/v2/internal/telemetry" + "google.golang.org/adk/v2/session" + "google.golang.org/adk/v2/session/compaction" +) + +// spanRecorder installs an in-memory tracer for the calling test. +// spanRecorder installs an in-memory tracer for the duration of a test. +// +// It replaces a package-level tracer, so no test in this file may call +// t.Parallel: a parallel test would swap the tracer while another test is +// reading it, which the race detector reports and which silently sends spans to +// the wrong exporter even when it does not. +func spanRecorder(t *testing.T) *tracetest.InMemoryExporter { + t.Helper() + exp := tracetest.NewInMemoryExporter() + telemetry.OverrideTracerForTesting(t, sdktrace.NewTracerProvider(sdktrace.WithSyncer(exp))) + return exp +} + +// attrs flattens a span's attributes for lookup by key. +func attrs(kvs []attribute.KeyValue) map[string]attribute.Value { + out := make(map[string]attribute.Value, len(kvs)) + for _, kv := range kvs { + out[string(kv.Key)] = kv.Value + } + return out +} + +func TestSlidingWindowEmitsSpan(t *testing.T) { + exp := spanRecorder(t) + + events := []*session.Event{ + textEvent("a", "inv1", 1, "q1"), modelTextEvent("b", "inv1", 2, "a1"), + textEvent("c", "inv2", 3, "q2"), modelTextEvent("d", "inv2", 4, "a2"), + } + cfg := &compaction.Config{CompactionInterval: 2, OverlapSize: 1, Summarizer: &fakeSummarizer{summary: "sum"}} + + got, err := slidingWindowStored(context.Background(), cfg, &staticSession{events: events}) + if err != nil { + t.Fatalf("slidingWindowStored() error = %v", err) + } + if got == nil { + t.Fatal("slidingWindowStored() produced no summary") + } + + spans := exp.GetSpans() + if len(spans) != 1 { + t.Fatalf("got %d spans, want 1", len(spans)) + } + span := spans[0] + if want := "compact_events sliding_window"; span.Name != want { + t.Errorf("span name = %q, want %q", span.Name, want) + } + if span.Status.Code == codes.Error { + t.Errorf("span status = error, want unset: %v", span.Status) + } + + a := attrs(span.Attributes) + for key, want := range map[string]string{ + "gen_ai.operation.name": "compact_events", + "gen_ai.conversation.id": "sess", + "gen_ai.compaction.trigger": "sliding_window", + "gen_ai.compaction.summarizer_type": "fakeSummarizer", + "gen_ai.compaction.result_event_id": got.ID, + } { + if a[key].AsString() != want { + t.Errorf("attribute %s = %q, want %q", key, a[key].AsString(), want) + } + } + if a["gen_ai.compaction.event_count"].AsInt64() != 4 { + t.Errorf("event_count = %d, want 4", a["gen_ai.compaction.event_count"].AsInt64()) + } + if a["gen_ai.compaction.compaction_interval"].AsInt64() != 2 { + t.Errorf("compaction_interval = %d, want 2", a["gen_ai.compaction.compaction_interval"].AsInt64()) + } + if a["gen_ai.compaction.overlap_size"].AsInt64() != 1 { + t.Errorf("overlap_size = %d, want 1", a["gen_ai.compaction.overlap_size"].AsInt64()) + } + // Only the knobs of the configured strategy appear. This says nothing + // about which strategy produced the span; the trigger attribute does. + if _, ok := a["gen_ai.compaction.token_threshold"]; ok { + t.Error("token_threshold attribute is present on a sliding-window span, want it omitted") + } + // The range must be recorded so a trace shows what the summary replaced, + // and it must be the right range in the right layout. Asserting only that + // the attributes are non-empty left the layout, the bound each one is + // sourced from, and the timestamps themselves all unprotected. + // Epoch seconds as a float, matching the reference implementation. The type + // is asserted as well as the value, because emitting these as strings is the + // defect this pins and a string attribute reads back as zero here. + wantStart := float64(at(1).UnixNano()) / float64(time.Second) + wantEnd := float64(at(4).UnixNano()) / float64(time.Second) + if got := a["gen_ai.compaction.start_timestamp"]; got.Type() != attribute.FLOAT64 || got.AsFloat64() != wantStart { + t.Errorf("start_timestamp = %v (%v), want %v (FLOAT64)", got.Emit(), got.Type(), wantStart) + } + if got := a["gen_ai.compaction.end_timestamp"]; got.Type() != attribute.FLOAT64 || got.AsFloat64() != wantEnd { + t.Errorf("end_timestamp = %v (%v), want %v (FLOAT64)", got.Emit(), got.Type(), wantEnd) + } + if got := a["gen_ai.compaction.result_event_id"].AsString(); got == "" { + t.Error("result_event_id is empty, so a trace cannot be joined to the stored summary") + } +} + +func TestCompactionSpanRecordsFailure(t *testing.T) { + exp := spanRecorder(t) + + events := []*session.Event{ + textEvent("a", "inv1", 1, "q1"), modelTextEvent("b", "inv1", 2, "a1"), + textEvent("c", "inv2", 3, "q2"), modelTextEvent("d", "inv2", 4, "a2"), + } + cfg := &compaction.Config{CompactionInterval: 2, Summarizer: &fakeSummarizer{err: errors.New("boom")}} + + if _, err := slidingWindowStored(context.Background(), cfg, &staticSession{events: events}); err == nil { + t.Fatal("slidingWindowStored() succeeded, want an error") + } + + spans := exp.GetSpans() + if len(spans) != 1 { + t.Fatalf("got %d spans, want 1", len(spans)) + } + if spans[0].Status.Code != codes.Error { + t.Errorf("span status = %v, want %v", spans[0].Status.Code, codes.Error) + } + if len(spans[0].Events) == 0 { + t.Error("span records no exception event, so the failure reason is lost") + } + // A failed compaction has no result. Recording one would leave a span that + // is at once an error and a success, naming an event nothing ever stored. + a := attrs(spans[0].Attributes) + for _, key := range []string{ + "gen_ai.compaction.result_event_id", + "gen_ai.compaction.start_timestamp", + "gen_ai.compaction.end_timestamp", + } { + if _, ok := a[key]; ok { + t.Errorf("%s is present on a failed compaction span, want it omitted", key) + } + } +} + +// TestNoSpanWhenNothingToCompact pins that evaluating a trigger and declining is +// silent, so the presence of a span in a trace means compaction really ran. +func TestNoSpanWhenNothingToCompact(t *testing.T) { + exp := spanRecorder(t) + + events := []*session.Event{textEvent("a", "inv1", 1, "q1")} + cfg := &compaction.Config{CompactionInterval: 5, Summarizer: &fakeSummarizer{summary: "sum"}} + + got, err := slidingWindowStored(context.Background(), cfg, &staticSession{events: events}) + if err != nil || got != nil { + t.Fatalf("slidingWindowStored() = (%v, %v), want (nil, nil)", got, err) + } + if n := len(exp.GetSpans()); n != 0 { + t.Errorf("got %d spans when the interval was not reached, want 0", n) + } +} + +// TestSpanRecordsDecliningSummarizer distinguishes "ran and produced nothing" +// from "ran and failed": the span exists and is successful, but carries no +// result attributes. +func TestSpanRecordsDecliningSummarizer(t *testing.T) { + exp := spanRecorder(t) + + events := []*session.Event{ + textEvent("a", "inv1", 1, "q1"), modelTextEvent("b", "inv1", 2, "a1"), + textEvent("c", "inv2", 3, "q2"), modelTextEvent("d", "inv2", 4, "a2"), + } + cfg := &compaction.Config{CompactionInterval: 2, Summarizer: &fakeSummarizer{}} + + if _, err := slidingWindowStored(context.Background(), cfg, &staticSession{events: events}); err != nil { + t.Fatalf("slidingWindowStored() error = %v", err) + } + + spans := exp.GetSpans() + if len(spans) != 1 { + t.Fatalf("got %d spans, want 1", len(spans)) + } + if spans[0].Status.Code == codes.Error { + t.Errorf("span status = error, want success for a summarizer that merely declined") + } + if _, ok := attrs(spans[0].Attributes)["gen_ai.compaction.result_event_id"]; ok { + t.Error("result_event_id is set although no summary was produced") + } +} + +// bothSummarizer returns a usable compaction event alongside an error, which a +// third-party Summarizer is free to do. +type bothSummarizer struct{} + +func (s *bothSummarizer) SummarizeEvents(_ context.Context, events []*session.Event) (compaction.SummarizeResult, error) { + // Content alongside an error. The framework must discard the content. + return compaction.SummarizeResult{Content: genai.NewContentFromText("SUM", "model")}, errors.New("boom") +} + +// TestCompactionSpanOmitsResultWhenSummarizerAlsoErrors pins that a span is +// never both an error and a success. +// +// A Summarizer may return an event and an error together. The caller discards +// the event, so recording its identity would name something no session holds, +// and the span would report a failure while carrying the attributes of a +// success. +func TestCompactionSpanOmitsResultWhenSummarizerAlsoErrors(t *testing.T) { + exp := spanRecorder(t) + + events := []*session.Event{ + textEvent("a", "inv1", 1, "q1"), modelTextEvent("b", "inv1", 2, "a1"), + textEvent("c", "inv2", 3, "q2"), modelTextEvent("d", "inv2", 4, "a2"), + } + cfg := &compaction.Config{CompactionInterval: 2, Summarizer: &bothSummarizer{}} + + if _, err := slidingWindowStored(context.Background(), cfg, &staticSession{events: events}); err == nil { + t.Fatal("slidingWindowStored() succeeded, want an error") + } + + spans := exp.GetSpans() + if len(spans) != 1 { + t.Fatalf("got %d spans, want 1", len(spans)) + } + if spans[0].Status.Code != codes.Error { + t.Errorf("span status = %v, want %v", spans[0].Status.Code, codes.Error) + } + a := attrs(spans[0].Attributes) + for _, key := range []string{ + "gen_ai.compaction.result_event_id", + "gen_ai.compaction.start_timestamp", + "gen_ai.compaction.end_timestamp", + } { + if v, ok := a[key]; ok { + t.Errorf("%s = %q on a failed compaction span, want it omitted", key, v.AsString()) + } + } +} + +// panickingSummarizer models third-party code that blows up. +type panickingSummarizer struct{} + +func (s *panickingSummarizer) SummarizeEvents(_ context.Context, _ []*session.Event) (compaction.SummarizeResult, error) { + panic("summarizer exploded") +} + +// TestCompactionSpanMarksAPanic pins that a panicking summarizer does not leave +// a span that reads as success. +// +// The OTel SDK records an exception event on the way out but leaves the status +// Unset, and Unset is indistinguishable from a healthy compaction that produced +// nothing. The panic itself still propagates. +func TestCompactionSpanMarksAPanic(t *testing.T) { + exp := spanRecorder(t) + + events := []*session.Event{ + textEvent("a", "inv1", 1, "q1"), modelTextEvent("b", "inv1", 2, "a1"), + textEvent("c", "inv2", 3, "q2"), modelTextEvent("d", "inv2", 4, "a2"), + } + cfg := &compaction.Config{CompactionInterval: 2, Summarizer: &panickingSummarizer{}} + + func() { + defer func() { + if r := recover(); r == nil { + t.Error("the panic did not propagate; compaction must not swallow it") + } + }() + _, _ = slidingWindowStored(context.Background(), cfg, &staticSession{events: events}) + }() + + spans := exp.GetSpans() + if len(spans) != 1 { + t.Fatalf("got %d spans, want 1", len(spans)) + } + if spans[0].Status.Code != codes.Error { + t.Errorf("span status = %v, want %v: a panicking summarizer must not look healthy", spans[0].Status.Code, codes.Error) + } +} + +// geminiSummarizer reports a backend the way the real summarizer does. +type geminiSummarizer struct { + fakeSummarizer + backend genai.Backend +} + +func (s *geminiSummarizer) GetGoogleLLMVariant() genai.Backend { return s.backend } + +// TestCompactionSpanRecordsGenAISystem pins gen_ai.system on the span. +// +// It names the system that produced the summary. Two deliberate divergences +// from the reference implementation, both repo-wide rather than compaction's, +// and both asserted here so a repo-wide change has to update this test rather +// than discover it in a dashboard. +// +// The values carry this repo's semconv prefix, "gcp.vertex_ai" against the +// reference's bare "vertex_ai". And a backend this repo cannot name leaves the +// attribute off, where the reference always emits one: naming a provider we +// have not identified is worse than saying nothing. Whatever the mapping is, it +// is shared with the rest of telemetry rather than restated here. +func TestCompactionSpanRecordsGenAISystem(t *testing.T) { + events := []*session.Event{ + textEvent("a", "inv1", 1, "q1"), modelTextEvent("b", "inv1", 2, "a1"), + textEvent("c", "inv2", 3, "q2"), modelTextEvent("d", "inv2", 4, "a2"), + } + + tests := []struct { + name string + backend genai.Backend + want string // "" means the attribute must be absent + }{ + {name: "vertex ai", backend: genai.BackendVertexAI, want: "gcp.vertex_ai"}, + {name: "gemini api", backend: genai.BackendGeminiAPI, want: "gcp.gemini"}, + {name: "summarizer that does not say", backend: genai.BackendUnspecified, want: ""}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + exp := spanRecorder(t) + cfg := &compaction.Config{ + CompactionInterval: 2, + Summarizer: &geminiSummarizer{ + fakeSummarizer: fakeSummarizer{summary: "SUM"}, + backend: tc.backend, + }, + } + if _, err := slidingWindowStored(context.Background(), cfg, &staticSession{events: events}); err != nil { + t.Fatalf("slidingWindowStored() error = %v", err) + } + spans := exp.GetSpans() + if len(spans) != 1 { + t.Fatalf("got %d spans, want 1", len(spans)) + } + // The expectation comes from the shared mapping, so this test + // tracks it rather than freezing a second copy of it. + if want, ok := telemetry.GenAISystemAttr(tc.backend); ok != (tc.want != "") || + (ok && want.Value.AsString() != tc.want) { + t.Fatalf("the shared mapping now returns (%v, %t) for %v, so this table is stale", want.Value.AsString(), ok, tc.backend) + } + got, ok := attrs(spans[0].Attributes)["gen_ai.system"] + if tc.want == "" { + if ok { + t.Errorf("gen_ai.system = %q, want it omitted", got.AsString()) + } + return + } + if !ok { + t.Fatal("gen_ai.system is absent") + } + if got.AsString() != tc.want { + t.Errorf("gen_ai.system = %q, want %q", got.AsString(), tc.want) + } + }) + } +} + +// TestCompactionSpanCarriesInvocationAndUsage pins the two attributes that let +// a compaction span be joined to the turn that caused it and costed. +// +// The span is not a child of the turn's span, so without the invocation id +// there is no way to ask which turn a compaction belonged to. And compaction +// spends a model call in order to save tokens later, so a span that does not +// record what it spent cannot show whether it paid for itself. +func TestCompactionSpanCarriesInvocationAndUsage(t *testing.T) { + exp := spanRecorder(t) + + events := []*session.Event{ + textEvent("a", "inv1", 1, "q1"), modelTextEvent("b", "inv1", 2, "a1"), + textEvent("c", "inv2", 3, "q2"), modelTextEvent("d", "inv2", 4, "a2"), + } + cfg := &compaction.Config{ + CompactionInterval: 2, + Summarizer: &usageSummarizer{ + fakeSummarizer: fakeSummarizer{summary: "SUM"}, + prompt: 1234, + output: 56, + }, + } + + if _, err := slidingWindowStored(context.Background(), cfg, &staticSession{events: events}); err != nil { + t.Fatalf("slidingWindowStored() error = %v", err) + } + a := attrs(exp.GetSpans()[0].Attributes) + + if got := a["gcp.vertex.agent.invocation_id"].AsString(); got != "inv2" { + t.Errorf("invocation_id = %q, want the turn that triggered compaction (%q)", got, "inv2") + } + if got := a["gen_ai.usage.input_tokens"].AsInt64(); got != 1234 { + t.Errorf("input_tokens = %d, want 1234", got) + } + if got := a["gen_ai.usage.output_tokens"].AsInt64(); got != 56 { + t.Errorf("output_tokens = %d, want 56", got) + } +} + +// usageSummarizer reports token usage the way a real one does. +type usageSummarizer struct { + fakeSummarizer + prompt int32 + output int32 +} + +func (s *usageSummarizer) SummarizeEvents(ctx context.Context, events []*session.Event) (compaction.SummarizeResult, error) { + res, err := s.fakeSummarizer.SummarizeEvents(ctx, events) + if err != nil || res.Content == nil { + return res, err + } + res.Usage = &genai.GenerateContentResponseUsageMetadata{ + PromptTokenCount: s.prompt, + CandidatesTokenCount: s.output, + } + return res, nil +} + +// TestCompactionSpanAttributeKeySet pins the exact set of attribute keys. +// +// The keys are a contract shared with adk-python, and the individual assertions +// elsewhere only check the keys they name. Adding, renaming or dropping one +// would otherwise pass unnoticed until a dashboard written against the other +// implementation stopped matching. +func TestCompactionSpanAttributeKeySet(t *testing.T) { + exp := spanRecorder(t) + + events := []*session.Event{ + textEvent("a", "inv1", 1, "q1"), modelTextEvent("b", "inv1", 2, "a1"), + textEvent("c", "inv2", 3, "q2"), modelTextEvent("d", "inv2", 4, "a2"), + } + cfg := &compaction.Config{CompactionInterval: 2, OverlapSize: 1, Summarizer: &fakeSummarizer{summary: "SUM"}} + + if _, err := slidingWindowStored(context.Background(), cfg, &staticSession{events: events}); err != nil { + t.Fatalf("slidingWindowStored() error = %v", err) + } + + want := []string{ + "gcp.vertex.agent.invocation_id", + "gen_ai.compaction.compaction_interval", + "gen_ai.compaction.end_timestamp", + "gen_ai.compaction.event_count", + "gen_ai.compaction.overlap_size", + "gen_ai.compaction.result_event_id", + "gen_ai.compaction.start_timestamp", + "gen_ai.compaction.summarizer_type", + "gen_ai.compaction.trigger", + "gen_ai.conversation.id", + "gen_ai.operation.name", + } + var got []string + for k := range attrs(exp.GetSpans()[0].Attributes) { + got = append(got, k) + } + slices.Sort(got) + if diff := cmp.Diff(want, got); diff != "" { + t.Errorf("attribute key set mismatch (-want +got):\n%s\nthese keys are shared with adk-python; change them together", diff) + } +} + +func TestTailRetentionEmitsSpan(t *testing.T) { + exp := spanRecorder(t) + + events := []*session.Event{ + textEvent("a", "inv1", 1, "q1"), + withUsage(modelTextEvent("b", "inv1", 2, "a1"), 900), + } + cfg := &compaction.Config{TokenThreshold: 100, EventRetentionSize: 0, Summarizer: &fakeSummarizer{summary: "sum"}} + + if _, err := tailRetentionStored(context.Background(), cfg, &staticSession{events: events}, TurnScope{}, nil, nil); err != nil { + t.Fatalf("tailRetentionStored() error = %v", err) + } + + spans := exp.GetSpans() + if len(spans) != 1 { + t.Fatalf("got %d spans, want 1", len(spans)) + } + if want := "compact_events token_threshold"; spans[0].Name != want { + t.Errorf("span name = %q, want %q", spans[0].Name, want) + } + a := attrs(spans[0].Attributes) + if a["gen_ai.compaction.token_threshold"].AsInt64() != 100 { + t.Errorf("token_threshold = %d, want 100", a["gen_ai.compaction.token_threshold"].AsInt64()) + } + if _, ok := a["gen_ai.compaction.compaction_interval"]; ok { + t.Error("compaction_interval attribute is present on a tail-retention span, want it omitted") + } +} + +// TestCompactionSpanRecordsTailRetentionThresholds pins the two attributes only +// a tail-retention span carries. +// +// They are declared in the telemetry commit, where nothing can exercise them +// because tail retention does not exist yet, so they were unprotected: renaming +// either one left the suite green. This is the first commit with a producer. +func TestCompactionSpanRecordsTailRetentionThresholds(t *testing.T) { + exp := spanRecorder(t) + + events := []*session.Event{ + textEvent("a", "inv1", 1, "q1"), modelTextEvent("b", "inv1", 2, "a1"), + textEvent("c", "inv2", 3, "q2"), modelTextEvent("d", "inv2", 4, "a2"), + } + cfg := &compaction.Config{ + TokenThreshold: 10, + EventRetentionSize: 1, + Summarizer: &fakeSummarizer{summary: "SUM"}, + } + + if _, err := tailRetentionStored(context.Background(), cfg, &staticSession{events: events}, TurnScope{}, func([]*session.Event) int { return 1000 }, nil); err != nil { + t.Fatalf("tailRetentionStored() error = %v", err) + } + + spans := exp.GetSpans() + if len(spans) != 1 { + t.Fatalf("got %d spans, want 1", len(spans)) + } + a := attrs(spans[0].Attributes) + if got := a["gen_ai.compaction.token_threshold"].AsInt64(); got != 10 { + t.Errorf("token_threshold = %d, want 10", got) + } + if got := a["gen_ai.compaction.event_retention_size"].AsInt64(); got != 1 { + t.Errorf("event_retention_size = %d, want 1", got) + } + // The knobs of the strategy that is not configured stay off the span. + if _, ok := a["gen_ai.compaction.compaction_interval"]; ok { + t.Error("compaction_interval is present on a tail-retention span, want it omitted") + } +} + +// TestCompactionSpanRecordsADecline pins the difference between a trigger that +// never fired and one that fired and could do nothing. +// +// The first stays silent, so a span in a trace still means compaction was +// wanted. The second used to be silent too, which made a session whose prompt +// grows on every turn look exactly like an idle one. +func TestCompactionSpanRecordsADecline(t *testing.T) { + exp := spanRecorder(t) + + // Threshold crossed, but the retained tail is the entire history, so there + // is nothing the compactor may summarize. + events := []*session.Event{ + textEvent("a", "inv1", 1, "q1"), modelTextEvent("b", "inv1", 2, "a1"), + } + cfg := &compaction.Config{ + TokenThreshold: 10, + EventRetentionSize: 50, + Summarizer: &fakeSummarizer{summary: "SUM"}, + } + + got, err := tailRetentionStored(context.Background(), cfg, &staticSession{events: events}, TurnScope{}, func([]*session.Event) int { return 1000 }, nil) + if err != nil || got != nil { + t.Fatalf("tailRetentionStored() = (%v, %v), want (nil, nil)", got, err) + } + + spans := exp.GetSpans() + if len(spans) != 1 { + t.Fatalf("got %d spans for a declined compaction, want 1", len(spans)) + } + reason, ok := attrs(spans[0].Attributes)["gen_ai.compaction.declined"] + if !ok { + t.Fatal("the span does not say it declined, so it is indistinguishable from one that compacted") + } + if reason.AsString() == "" { + t.Error("the decline reason is empty") + } + if n := attrs(spans[0].Attributes)["gen_ai.compaction.event_count"].AsInt64(); n != 0 { + t.Errorf("event_count = %d on a declined compaction, want 0", n) + } +} + +// TestCompactionSpanOmitsAbsentTimestamps pins that a range with no bounds is +// reported as absent rather than as the year 1754. +// +// Epoch seconds of a zero time is -6.795e+09, so a compaction covering three +// seconds of history was published as a range 271 years wide, on a span that +// otherwise reported success. An absent attribute is the only form a consumer +// can tell apart from a real reading. +// +// The range is derived from the covered events, so the way to reach an unset +// bound is a covered event that was never stamped, which is what these events +// are. Nothing on the append path used to fill a missing timestamp in. +func TestCompactionSpanOmitsAbsentTimestamps(t *testing.T) { + exp := spanRecorder(t) + + // Only the oldest event is unstamped, so the window still has a real upper + // bound and the selection logic, which compares against the previous + // compaction's end, still sees the later invocations as new. + first := textEvent("a", "inv1", 1, "q1") + first.Timestamp = time.Time{} + events := []*session.Event{ + first, modelTextEvent("b", "inv1", 2, "a1"), + textEvent("c", "inv2", 3, "q2"), modelTextEvent("d", "inv2", 4, "a2"), + } + cfg := &compaction.Config{CompactionInterval: 2, Summarizer: &fakeSummarizer{summary: "SUM"}} + + if _, err := slidingWindowStored(context.Background(), cfg, &staticSession{events: events}); err != nil { + t.Fatalf("slidingWindowStored() error = %v", err) + } + a := attrs(exp.GetSpans()[0].Attributes) + + if v, ok := a["gen_ai.compaction.start_timestamp"]; ok { + t.Errorf("start_timestamp = %v, want the key to be absent for an unset bound", v.AsFloat64()) + } + // The bound that does exist is still reported, so absence means absence + // rather than the attribute simply being dropped. + if _, ok := a["gen_ai.compaction.end_timestamp"]; !ok { + t.Error("end_timestamp is missing, want the bound that was recorded") + } +} + +// TestCompactionSpanReportsADiscardedSummary pins that a summary the caller +// threw away is not reported as a stored one. +// +// The span used to end when the summarizer returned, before any of the reasons +// a caller discards a result: a cancelled turn, a failed re-read, a competing +// compaction, a plugin rejecting it, a failed append. Every one of those left a +// span saying the compaction succeeded, carrying a result_event_id that exists +// in no session, so a trace could not distinguish a compaction that shrank a +// prompt from one that spent a model call and changed nothing. +func TestCompactionSpanReportsADiscardedSummary(t *testing.T) { + exp := spanRecorder(t) + + events := []*session.Event{ + textEvent("a", "inv1", 1, "q1"), modelTextEvent("b", "inv1", 2, "a1"), + textEvent("c", "inv2", 3, "q2"), modelTextEvent("d", "inv2", 4, "a2"), + } + cfg := &compaction.Config{CompactionInterval: 2, Summarizer: &fakeSummarizer{summary: "SUM"}} + + summary, finish, err := SlidingWindow(context.Background(), cfg, &staticSession{events: events}, "") + if err != nil || summary == nil { + t.Fatalf("SlidingWindow() = %v, %v, want a summary", summary, err) + } + // The caller decides not to keep it. + finish(nil, "another compaction covering the same events landed while summarizing") + + spans := exp.GetSpans() + if len(spans) != 1 { + t.Fatalf("recorded %d spans, want 1", len(spans)) + } + a := attrs(spans[0].Attributes) + if got, ok := a["gen_ai.compaction.declined"]; !ok { + t.Error("the span does not say the summary was discarded") + } else if got.AsString() == "" { + t.Error("the discard reason is empty") + } + if _, ok := a["gen_ai.compaction.result_event_id"]; ok { + t.Error("the span names a result event, but nothing reached the session") + } +} + +// TestCompactionSpanRecordsAPanicAsAnException pins that a panicking summarizer +// is visible to an alert keyed on exception.type. +// +// The status was set to Error, which a dashboard sees, but no exception event +// was recorded, so the panic itself was invisible to the usual alerting path. +func TestCompactionSpanRecordsAPanicAsAnException(t *testing.T) { + exp := spanRecorder(t) + + events := []*session.Event{ + textEvent("a", "inv1", 1, "q1"), modelTextEvent("b", "inv1", 2, "a1"), + textEvent("c", "inv2", 3, "q2"), modelTextEvent("d", "inv2", 4, "a2"), + } + cfg := &compaction.Config{CompactionInterval: 2, Summarizer: &panickingSummarizer{}} + + func() { + defer func() { _ = recover() }() + _, _, _ = SlidingWindow(context.Background(), cfg, &staticSession{events: events}, "") + }() + + spans := exp.GetSpans() + if len(spans) != 1 { + t.Fatalf("recorded %d spans, want 1", len(spans)) + } + var sawException bool + for _, e := range spans[0].Events { + if e.Name == "exception" { + sawException = true + } + } + if !sawException { + t.Error("no exception event was recorded, so an alert keyed on exception.type misses a panicking summarizer") + } +} + +// TestDeclinedSummarizationReportsWhatItSpent pins that a decline is visible on +// the span, with the usage the summarizer reported. +// +// Three places disagreed about this and the code followed the one that did +// nothing: the interface invites a summarizer to report usage alongside a +// decline, the compactor's comment said the span records it, and the span was +// left bare. So a decline looked identical to a compaction that was never +// attempted, and a summarizer that spent a model call and got nothing usable +// back reported no spend at all. +func TestDeclinedSummarizationReportsWhatItSpent(t *testing.T) { + // Not parallel: spanRecorder overrides the global tracer, so every test in + // this file that records spans has to run alone. It is the only one here + // that was. + + exp := spanRecorder(t) + + events := []*session.Event{ + textEvent("a", "inv1", 1, "q1"), + modelTextEvent("b", "inv1", 2, "a1"), + textEvent("c", "inv2", 3, "q2"), + modelTextEvent("d", "inv2", 4, "a2"), + } + cfg := &compaction.Config{ + CompactionInterval: 2, OverlapSize: 1, + Summarizer: &usageReportingDecliner{usage: &genai.GenerateContentResponseUsageMetadata{ + PromptTokenCount: 120, CandidatesTokenCount: 15, + }}, + } + if _, err := slidingWindowStored(context.Background(), cfg, &staticSession{events: events}); err != nil { + t.Fatalf("slidingWindowStored() error = %v", err) + } + + spans := exp.GetSpans() + if len(spans) != 1 { + t.Fatalf("recorded %d spans, want 1", len(spans)) + } + got := attrs(spans[0].Attributes) + if _, ok := got["gen_ai.compaction.declined"]; !ok { + t.Error("a decline left no marker on the span, so it is indistinguishable from compaction never being attempted") + } + if _, ok := got["gen_ai.usage.input_tokens"]; !ok { + t.Error("the span records no input tokens for a decline that reported usage") + } + if _, ok := got["gen_ai.usage.output_tokens"]; !ok { + t.Error("the span records no output tokens for a decline that reported usage") + } +} + +// usageReportingDecliner declines while reporting what the attempt cost. +type usageReportingDecliner struct { + usage *genai.GenerateContentResponseUsageMetadata +} + +func (d *usageReportingDecliner) SummarizeEvents(context.Context, []*session.Event) (compaction.SummarizeResult, error) { + return compaction.SummarizeResult{Usage: d.usage}, nil +} + +// TestFailedSummarizationReportsWhatItSpent pins the same property for a +// failure, which is the case that costs the most. +// +// The built-in summarizer returns usage alongside the error it raises for a +// MAX_TOKENS or safety stop, having already paid for a full-transcript prompt. +// Usage was recorded after the early return taken for an error, so a session +// failing every compaction reported no spend at all while burning the largest +// prompts it ever sends. +func TestFailedSummarizationReportsWhatItSpent(t *testing.T) { + // Not parallel: spanRecorder overrides the global tracer. + + exp := spanRecorder(t) + + events := []*session.Event{ + textEvent("a", "inv1", 1, "q1"), + modelTextEvent("b", "inv1", 2, "a1"), + textEvent("c", "inv2", 3, "q2"), + modelTextEvent("d", "inv2", 4, "a2"), + } + cfg := &compaction.Config{ + CompactionInterval: 2, OverlapSize: 1, + Summarizer: &usageReportingFailer{usage: &genai.GenerateContentResponseUsageMetadata{ + PromptTokenCount: 4000, CandidatesTokenCount: 8, + }}, + } + if _, err := slidingWindowStored(context.Background(), cfg, &staticSession{events: events}); err == nil { + t.Fatal("slidingWindowStored() error = nil, want the summarizer's failure") + } + + spans := exp.GetSpans() + if len(spans) != 1 { + t.Fatalf("recorded %d spans, want 1", len(spans)) + } + got := attrs(spans[0].Attributes) + if _, ok := got["gen_ai.usage.input_tokens"]; !ok { + t.Error("the span records no input tokens for a failure that reported usage, so the spend is invisible") + } + if _, ok := got["gen_ai.usage.output_tokens"]; !ok { + t.Error("the span records no output tokens for a failure that reported usage") + } + if _, ok := got["gen_ai.compaction.result_event_id"]; ok { + t.Error("a failed compaction named a result event, so the span is at once an error and a success") + } +} + +// usageReportingFailer fails while reporting what the attempt cost, as the +// built-in summarizer does for a MAX_TOKENS or safety stop. +type usageReportingFailer struct { + usage *genai.GenerateContentResponseUsageMetadata +} + +func (f *usageReportingFailer) SummarizeEvents(context.Context, []*session.Event) (compaction.SummarizeResult, error) { + return compaction.SummarizeResult{Usage: f.usage}, errors.New(`summarizer stopped before finishing (finish reason "MAX_TOKENS"), so the summary is incomplete`) +} diff --git a/internal/llminternal/base_flow.go b/internal/llminternal/base_flow.go index a5d03bc6f..e33eeb33a 100644 --- a/internal/llminternal/base_flow.go +++ b/internal/llminternal/base_flow.go @@ -83,6 +83,9 @@ var ( RequestConfirmationRequestProcessor, instructionsRequestProcessor, identityRequestProcessor, + // Compaction must run before contentsRequestProcessor so a summary it + // appends is reflected in the history assembled for this very request. + CompactionRequestProcessor, ContentsRequestProcessor, // Some implementations of NL Planning mark planning contents as thoughts in the post processor. // Since these need to be unmarked, NL Planning should be after contentsRequestProcessor. diff --git a/internal/llminternal/compaction_processor.go b/internal/llminternal/compaction_processor.go new file mode 100644 index 000000000..911e77fd7 --- /dev/null +++ b/internal/llminternal/compaction_processor.go @@ -0,0 +1,247 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package llminternal + +import ( + "context" + "fmt" + "iter" + "log" + + "google.golang.org/adk/v2/agent" + "google.golang.org/adk/v2/internal/agent/compactionctx" + "google.golang.org/adk/v2/internal/compactioninternal" + "google.golang.org/adk/v2/model" + "google.golang.org/adk/v2/session" + "google.golang.org/adk/v2/session/compaction" +) + +// CompactionRequestProcessor runs token-threshold tail-retention compaction +// before the conversation history is assembled for a model call. +// +// It must sit before [ContentsRequestProcessor] in the chain: the summary it +// appends only shrinks this request if contents are built afterwards. +// +// Unlike the runner's post-invocation sliding-window pass, this runs mid-turn, +// so it can react to a single long-running invocation that inflates the prompt +// rather than waiting for the turn to finish. It leaves the request itself +// untouched and emits no events. +func CompactionRequestProcessor(ctx agent.InvocationContext, _ *model.LLMRequest, _ *Flow) iter.Seq2[*session.Event, error] { + return func(yield func(*session.Event, error) bool) { + rt := compactionctx.FromContext(ctx) + if !rt.Enabled() { + return + } + if ctx.Session() == nil { + return + } + + // Compact against the session underneath any wrapper an agent installed + // over it. A wrapper carries a synthetic first-turn seed that no store + // holds, and every session service type-asserts on its own concrete + // type, so appending a summary to one fails outright. + // + // Unwrapping rather than re-reading keeps object identity with the + // session the wrapper reads through, so the summary appended below + // reaches the prompt this processor runs ahead of. A freshly read + // session would be a different object and the summary would miss it. + sess := compactioninternal.UnwrapSession(ctx.Session()) + + // Which events exist before the model call, captured now rather than + // read back later. sess is the live handle and every backend mutates + // it in place, so passing it as the "before" state to the race check + // below compares the session against itself: anything a sibling + // appended during the call is already in it, the check finds no + // difference, and the record is stored claiming to cover an event no + // summary describes. Sub-agents make that ordinary rather than + // theoretical, because parallel and workflow nodes hand this very + // session down to each child. + before := compactioninternal.KnownEventIDs(sess) + + // Compaction is an optimisation, so a cancelled or expired turn should + // not spend a model call on it. + if ctx.Err() != nil { + return + } + summary, finish, err := compactioninternal.TailRetention(ctx, rt.Config(), sess, compactioninternal.TurnScope{ + InvocationID: ctx.InvocationID(), + Branch: ctx.Branch(), + IsolationScope: ctx.IsolationScope(), + }, promptTokenEstimator(ctx), rt.GateFor(ctx.Agent().Name(), ctx.Branch(), ctx.IsolationScope())) + if err != nil { + degrade(ctx, "token-threshold", err) + return + } + if summary == nil { + return + } + + // Summarizing takes a model call, which is long enough for another + // invocation on this session to append inside the range just chosen. + // Read the stored session and abandon the summary if anything landed + // inside it. Skipping costs one wasted call, where recording it would + // silently drop those turns from every later prompt. + // + // The read is only a comparison. The append below still goes to sess, + // for the identity reason above. + if ctx.Err() != nil { + finish(nil, "the turn ended before the summary could be stored") + return + } + latest, err := compactioninternal.ReloadSession(ctx, rt.SessionService(), sess) + if err != nil { + // Same reasoning as a failed summarization: this is bookkeeping in + // the middle of a turn whose tools may already have run. Failing to + // re-read means we cannot prove the summary is safe to keep, so it + // is dropped, but the turn continues. + finish(err, "") + degrade(ctx, "token-threshold", err) + return + } + if compactioninternal.RangeRacedSince(latest, before, summary) { + finish(nil, "another compaction covering the same events landed while summarizing") + log.Printf("adk: discarding a tail-retention summary because the session changed inside its range while summarizing") + return + } + + // Re-checked before it is stored, the way the sliding-window path + // re-checks a summary a plugin has touched. Nothing rewrites this one + // today, so this is a guard rather than a fix, and it is here because + // the two paths should not disagree about what may reach a prompt. + // + // The plugin pipeline itself is still not run here. A redaction plugin + // therefore sees every sliding-window summary and none of these, which + // is a real gap and a behaviour change to close rather than a bug to + // patch quietly. ADK Kotlin, which this design was adapted from, runs + // no plugin hook on either of its compaction paths, so the gap is + // consistent with the reference; that is a statement about consistency + // rather than a defence of the behaviour. It is documented on the + // exported surface at compaction.Config.TokenThreshold. + if !compactioninternal.SanitizeSummary(summary) { + finish(nil, "the summary held nothing usable") + log.Printf("adk: discarding a tail-retention summary because it held no usable content") + return + } + + if err := rt.SessionService().AppendEvent(ctx, sess, summary); err != nil { + finish(err, "") + degrade(ctx, "failed to append the summary event", err) + return + } + finish(nil, "") + + // The append itself cannot be guarded, so what lands during it is + // repaired afterwards, the same way the post-invocation path does. + // + // This path needs it more, not less. It runs inside the invocation + // rather than after it, so sub-agents and tools are still producing + // events while the summary is being stored, and tail retention is the + // strategy an application relies on for a bound. An event stored + // between the race check and the append sits inside the recorded range + // named by nothing, and prompt assembly drops it from every later + // prompt with no summary standing in for it. + // + // A failure here leaves the summary stored and correct for everything + // except a straggler, which is where this path was before the repair + // existed, so it is logged rather than surfaced: the turn is mid-flight + // and its tools may already have run. + // + // Detached from the caller's cancellation: the summary is already + // stored and claims a range wider than it summarized until this lands. + repairCtx, cancelRepair := compactioninternal.RepairContext(ctx) + defer cancelRepair() + if latest, err := compactioninternal.ReloadSession(repairCtx, rt.SessionService(), sess); err != nil { + log.Printf("adk: could not re-read the session to check a stored compaction for stragglers: %v", err) + } else if repair := compactioninternal.RepairAfterAppend(summary, before, latest); repair != nil { + if err := rt.SessionService().AppendEvent(repairCtx, sess, repair); err != nil { + log.Printf("adk: could not store a corrected compaction record: %v", err) + } else { + log.Printf("adk: corrected a tail-retention record that would have covered %d event(s) it did not summarize", + len(repair.Actions.Compaction.ExcludedEvents)-len(summary.Actions.Compaction.ExcludedEvents)) + } + } + + // The post-invocation sliding window checks this and stands down, so a + // turn that was compacted mid-flight is not summarized twice. + rt.MarkCompacted() + } +} + +// degrade reports a failed mid-turn compaction and lets the turn continue. +// +// This runs before a model call, in the middle of an invocation whose tools may +// already have run and committed their side effects. Failing the turn for a +// failed optimisation is never the right trade there: the user loses an answer, +// the side effects stand, and any summary already written is orphaned. Letting +// it through costs a larger prompt, and the model call either succeeds anyway, +// because the threshold sits well below the real context limit, or fails with +// the provider's own error, which says more about the actual problem than a +// compaction error would. +// +// The failure is not lost. It is logged, and the compaction span records it +// with an error status, so a summarizer failing every call is visible in traces +// rather than only in an aborted turn. The post-invocation pass still surfaces +// its own failures to the caller, since nothing is mid-flight there. +func degrade(ctx context.Context, stage string, err error) { + log.Printf("adk: %v; continuing with an uncompacted prompt", compactionFailure(stage, err)) +} + +// compactionFailure marks err as a compaction failure at the named stage. +// +// The cause is rendered with %v rather than wrapped with %w, deliberately. This +// error is yielded into the flow's error channel, which reaches the workflow +// scheduler, and the scheduler tests for a context.Canceled chain before +// anything else and drops the error when it finds one. A summarizer that failed +// because its own context was cancelled would therefore end the turn with no +// answer, no events and no error at all: the most confusing outcome available. +// +// Cutting the chain keeps the cause in the message and keeps the error matchable +// as [compaction.ErrCompaction], at the cost of errors.Is against the cause. For +// a bookkeeping failure that is the right way round. +func compactionFailure(stage string, err error) error { + return fmt.Errorf("%w: %s: %v", compaction.ErrCompaction, stage, err) +} + +// promptTokenEstimator returns a [compactioninternal.TokenCounter] that approximates +// the prompt size for ctx's agent. +// +// It is only consulted before any model response has reported a real token +// count. Building the contents the same way the request will is what makes the +// estimate meaningful: it sees branch and isolation-scope filtering, and any +// compaction already applied. +func promptTokenEstimator(ctx agent.InvocationContext) compactioninternal.TokenCounter { + return func(events []*session.Event) int { + llmAgent := asLLMAgent(ctx.Agent()) + if llmAgent == nil { + return 0 + } + state := llmAgent.internal() + contents, err := buildContentsDefault( + ctx.Agent().Name(), + ctx.Branch(), + ctx.IsolationScope(), + events, + state.Mode == ModeSingleTurn, + ctx.UserContent(), + ) + if err != nil { + // An unbuildable history is the contents processor's problem to + // report a moment from now; here it just means "no estimate". + return 0 + } + + return compactioninternal.EstimateTokensFromContents(contents) + } +} diff --git a/internal/llminternal/compaction_processor_test.go b/internal/llminternal/compaction_processor_test.go new file mode 100644 index 000000000..d6aca3673 --- /dev/null +++ b/internal/llminternal/compaction_processor_test.go @@ -0,0 +1,327 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package llminternal_test + +import ( + "context" + "testing" + "time" + + "google.golang.org/genai" + + "google.golang.org/adk/v2/agent/llmagent" + "google.golang.org/adk/v2/internal/agent/compactionctx" + "google.golang.org/adk/v2/internal/compactioninternal" + icontext "google.golang.org/adk/v2/internal/context" + "google.golang.org/adk/v2/internal/llminternal" + "google.golang.org/adk/v2/internal/utils" + "google.golang.org/adk/v2/model" + "google.golang.org/adk/v2/session" + "google.golang.org/adk/v2/session/compaction" +) + +// seedWrappedSession stands in for the wrapper an agent installs over the real +// session when it hands a sub-agent a synthetic first turn. It decorates the +// session and is not a type any session service recognises. +type seedWrappedSession struct { + session.Session +} + +func (w *seedWrappedSession) Unwrap() session.Session { return w.Session } + +// fixedSummarizer returns one canned summary. +type fixedSummarizer struct{ calls int } + +func (s *fixedSummarizer) SummarizeEvents(_ context.Context, events []*session.Event) (compaction.SummarizeResult, error) { + s.calls++ + return compaction.SummarizeResult{Content: genai.NewContentFromText("SUMMARY", "model")}, nil +} + +// tailRetentionFixture builds a stored session holding n exchanges, the last of +// which reports a prompt token count well past any threshold a test will set. +func tailRetentionFixture(t *testing.T, n int) (session.Service, session.Session) { + t.Helper() + + svc := session.InMemoryService() + created, err := svc.Create(t.Context(), &session.CreateRequest{ + AppName: "app", UserID: "u", SessionID: "s", + }) + if err != nil { + t.Fatalf("Create() error = %v", err) + } + sess := created.Session + + base := time.Unix(1, 0) + for i := range n { + q := session.NewEvent(t.Context(), "inv") + q.Author = "user" + q.Timestamp = base.Add(time.Duration(2*i) * time.Second) + q.LLMResponse.Content = genai.NewContentFromText("question", "user") + if err := svc.AppendEvent(t.Context(), sess, q); err != nil { + t.Fatalf("AppendEvent() error = %v", err) + } + + a := session.NewEvent(t.Context(), "inv") + a.Author = "assistant" + a.Timestamp = base.Add(time.Duration(2*i+1) * time.Second) + a.LLMResponse.Content = genai.NewContentFromText("answer", "model") + a.LLMResponse.UsageMetadata = &genai.GenerateContentResponseUsageMetadata{PromptTokenCount: 5000} + if err := svc.AppendEvent(t.Context(), sess, a); err != nil { + t.Fatalf("AppendEvent() error = %v", err) + } + } + return svc, sess +} + +func runCompactionProcessor(t *testing.T, svc session.Service, sess session.Session, cfg *compaction.Config) error { + t.Helper() + + ctx := compactionctx.ToContext(t.Context(), compactionctx.New(cfg, svc)) + testAgent := utils.Must(llmagent.New(llmagent.Config{Name: "assistant", Model: &testModel{}})) + ictx := icontext.NewInvocationContext(ctx, icontext.InvocationContextParams{ + Agent: testAgent, + Session: sess, + }) + + var gotErr error + for ev, err := range llminternal.CompactionRequestProcessor(ictx, &model.LLMRequest{}, &llminternal.Flow{}) { + if ev != nil { + t.Fatal("CompactionRequestProcessor yielded an event, which it must not do") + } + if err != nil { + gotErr = err + } + } + return gotErr +} + +func storedCompactions(t *testing.T, svc session.Service) []*session.Event { + t.Helper() + + resp, err := svc.Get(t.Context(), &session.GetRequest{AppName: "app", UserID: "u", SessionID: "s"}) + if err != nil { + t.Fatalf("Get() error = %v", err) + } + var out []*session.Event + for ev := range resp.Session.Events().All() { + if compactioninternal.HasUsableSummary(ev) { + out = append(out, ev) + } + } + return out +} + +// TestCompactionProcessorAppendsThroughASessionWrapper checks that tail +// retention still works when the invocation carries a wrapped session. +// +// An agent hands a sub-agent a session wrapped to carry a synthetic first turn. +// Every session service type-asserts on its own concrete type, so appending a +// summary to the wrapper fails outright. The failure does not even surface as an +// error on the delegating path: it becomes a tool-error response and the +// coordinator answers on top of a broken delegation. +func TestCompactionProcessorAppendsThroughASessionWrapper(t *testing.T) { + t.Parallel() + + svc, sess := tailRetentionFixture(t, 4) + summarizer := &fixedSummarizer{} + + err := runCompactionProcessor(t, svc, &seedWrappedSession{Session: sess}, &compaction.Config{ + TokenThreshold: 100, + EventRetentionSize: 2, + Summarizer: summarizer, + }) + if err != nil { + t.Fatalf("compaction through a wrapped session failed: %v", err) + } + if summarizer.calls == 0 { + t.Fatal("the summarizer never ran, so this test proved nothing") + } + if got := len(storedCompactions(t, svc)); got != 1 { + t.Errorf("stored %d compaction events, want 1: the summary never reached the session", got) + } +} + +// TestCompactionProcessorSkipsOnCancelledContext checks that a cancelled turn +// does not spend a model call on compaction, nor write a summary the caller is +// no longer waiting for. +func TestCompactionProcessorSkipsOnCancelledContext(t *testing.T) { + t.Parallel() + + svc, sess := tailRetentionFixture(t, 4) + summarizer := &fixedSummarizer{} + + ctx, cancel := context.WithCancel(t.Context()) + ctx = compactionctx.ToContext(ctx, compactionctx.New(&compaction.Config{ + TokenThreshold: 100, + EventRetentionSize: 2, + Summarizer: summarizer, + }, svc)) + testAgent := utils.Must(llmagent.New(llmagent.Config{Name: "assistant", Model: &testModel{}})) + ictx := icontext.NewInvocationContext(ctx, icontext.InvocationContextParams{ + Agent: testAgent, + Session: sess, + }) + cancel() + + for range llminternal.CompactionRequestProcessor(ictx, &model.LLMRequest{}, &llminternal.Flow{}) { //nolint:revive + } + + if summarizer.calls != 0 { + t.Errorf("summarizer ran %d time(s) on a cancelled turn", summarizer.calls) + } + if got := len(storedCompactions(t, svc)); got != 0 { + t.Errorf("stored %d compaction events on a cancelled turn", got) + } +} + +// racingSummarizer appends an event inside the range it is about to summarize, +// standing in for a concurrent invocation landing during the model call. +type racingSummarizer struct { + svc session.Service + t *testing.T + + // The event this summarizer landed mid-call. + racedInvocation string + racedTimestamp time.Time +} + +func (s *racingSummarizer) SummarizeEvents(ctx context.Context, events []*session.Event) (compaction.SummarizeResult, error) { + // Land a new event inside the range the summary claims, through a separate + // handle on the same stored session. Appending through the caller's handle + // would update that handle too, which is exactly what a concurrent + // invocation in another goroutine does not do. + other, err := s.svc.Get(ctx, &session.GetRequest{AppName: "app", UserID: "u", SessionID: "s"}) + if err != nil { + s.t.Fatalf("racing Get() error = %v", err) + } + late := session.NewEvent(ctx, "other-invocation") + late.Author = "user" + // Inside the range the framework will derive from the window it handed over. + late.Timestamp = events[0].Timestamp.Add(time.Millisecond) + late.LLMResponse.Content = genai.NewContentFromText("CONCURRENT", "user") + if err := s.svc.AppendEvent(ctx, other.Session, late); err != nil { + s.t.Fatalf("racing AppendEvent() error = %v", err) + } + s.racedInvocation, s.racedTimestamp = late.InvocationID, late.Timestamp + return compaction.SummarizeResult{Content: genai.NewContentFromText("SUMMARY", "model")}, nil +} + +// TestCompactionProcessorDiscardsARacedSummary checks that a summary is thrown +// away when another invocation appended inside its range while it was being +// produced. +// +// A summary records the holes inside its range, and that list is computed from +// what the framework could see when it was built. An event that lands +// afterwards is inside the range and named by nothing, so it reads as covered +// and prompt assembly drops it, having been summarized by nothing. Discarding +// costs one wasted model call; keeping it costs a turn of conversation. +func TestCompactionProcessorDiscardsARacedSummary(t *testing.T) { + t.Parallel() + + svc, sess := tailRetentionFixture(t, 4) + + summarizer := &racingSummarizer{svc: svc, t: t} + err := runCompactionProcessor(t, svc, sess, &compaction.Config{ + TokenThreshold: 100, + EventRetentionSize: 2, + Summarizer: summarizer, + }) + if err != nil { + t.Fatalf("CompactionRequestProcessor failed: %v", err) + } + if summarizer.racedInvocation == "" { + t.Fatal("the racing summarizer did not record what it appended") + } + if got := len(storedCompactions(t, svc)); got != 0 { + t.Errorf("stored %d compaction events, want 0: a summary whose range was raced must be discarded", got) + } +} + +// straggleOnAppend lands an event inside the summary's range at the moment the +// summary itself is being stored, which is the one window the race guard cannot +// cover. +type straggleOnAppend struct { + session.Service + t *testing.T + done bool + + stragglerInvocation string + stragglerTimestamp time.Time +} + +func (s *straggleOnAppend) AppendEvent(ctx context.Context, sess session.Session, ev *session.Event) error { + if !s.done && ev.Actions.Compaction != nil { + s.done = true + // Through a separate handle, the way a concurrent invocation would. + other, err := s.Service.Get(ctx, &session.GetRequest{AppName: "app", UserID: "u", SessionID: "s"}) + if err != nil { + s.t.Fatalf("straggler Get() error = %v", err) + } + late := session.NewEvent(ctx, "straggler-invocation") + late.Author = "user" + late.Timestamp = ev.Actions.Compaction.StartTimestamp.Add(time.Millisecond) + late.LLMResponse.Content = genai.NewContentFromText("PLEASE DO NOT LOSE ME", "user") + if err := s.Service.AppendEvent(ctx, other.Session, late); err != nil { + s.t.Fatalf("straggler AppendEvent() error = %v", err) + } + s.stragglerInvocation, s.stragglerTimestamp = late.InvocationID, late.Timestamp + } + return s.Service.AppendEvent(ctx, sess, ev) +} + +// TestCompactionProcessorRepairsAStragglerStoredDuringTheAppend pins that the +// mid-turn path repairs what the race guard cannot cover. +// +// The guard runs immediately before the append and the append cannot be made +// conditional, so an event stored in between lands inside the recorded range +// named by no hole, and prompt assembly drops it with no summary standing in +// for it. The post-invocation path has repaired this since it was found; this +// path did not, and it is the one that needs it more, because it runs inside +// the invocation while sub-agents and tools are still producing events. +func TestCompactionProcessorRepairsAStragglerStoredDuringTheAppend(t *testing.T) { + t.Parallel() + + base, sess := tailRetentionFixture(t, 6) + svc := &straggleOnAppend{Service: base, t: t} + + if err := runCompactionProcessor(t, svc, sess, &compaction.Config{ + TokenThreshold: 100, + EventRetentionSize: 2, + Summarizer: &fixedSummarizer{}, + }); err != nil { + t.Fatalf("CompactionRequestProcessor failed: %v", err) + } + if svc.stragglerInvocation == "" { + t.Fatal("no straggler was landed, so this test is not exercising the window") + } + + got, err := svc.Get(t.Context(), &session.GetRequest{AppName: "app", UserID: "u", SessionID: "s"}) + if err != nil { + t.Fatalf("Get() error = %v", err) + } + var stored []*session.Event + for ev := range got.Session.Events().All() { + stored = append(stored, ev) + } + + // The straggler must survive into the assembled prompt: no summary + // describes it, so it has to be there raw. + for _, ev := range compactioninternal.Apply(stored) { + if ev.InvocationID == svc.stragglerInvocation { + return + } + } + t.Errorf("the event stored during the append was dropped from the prompt and no summary describes it") +} diff --git a/runner/compaction_test.go b/runner/compaction_test.go index 968c1d7f6..4916ec587 100644 --- a/runner/compaction_test.go +++ b/runner/compaction_test.go @@ -1011,6 +1011,395 @@ func TestCompactionSpanJoinsTheCallersTrace(t *testing.T) { // usageModel replies with a canned answer and reports a fixed prompt token // count, so tail-retention compaction can be driven deterministically. +type usageModel struct { + mu sync.Mutex + prompts [][]*genai.Content + promptTokens int32 +} + +func (m *usageModel) Name() string { return "usage" } + +func (m *usageModel) GenerateContent(_ context.Context, req *model.LLMRequest, _ bool) iter.Seq2[*model.LLMResponse, error] { + m.mu.Lock() + m.prompts = append(m.prompts, req.Contents) + n := len(m.prompts) + tokens := m.promptTokens + m.mu.Unlock() + + return func(yield func(*model.LLMResponse, error) bool) { + yield(&model.LLMResponse{ + Content: genai.NewContentFromText(fmt.Sprintf("answer %d", n), "model"), + UsageMetadata: &genai.GenerateContentResponseUsageMetadata{PromptTokenCount: tokens}, + }, nil) + } +} + +func (m *usageModel) lastPrompt() []*genai.Content { + m.mu.Lock() + defer m.mu.Unlock() + if len(m.prompts) == 0 { + return nil + } + return m.prompts[len(m.prompts)-1] +} + +func TestRunnerTailRetentionCompactsMidInvocation(t *testing.T) { + t.Parallel() + + const userID, sessionID = "u", "s" + // Every model call reports a prompt well past the threshold, so compaction + // fires as soon as there are more events than the retained tail. + m := &usageModel{promptTokens: 5000} + summarizer := &recordingSummarizer{summary: "TAIL-SUMMARY"} + // Retention 1, because the question that opens the turn being answered is + // held back on top of the retained tail rather than counting towards it. + // At retention 2 the three events of this fixture are all spoken for. + r, svc := newCompactionRunner(t, m, &compaction.Config{ + TokenThreshold: 1000, + EventRetentionSize: 1, + Summarizer: summarizer, + }) + + // First turn: no prior usage metadata and only the user event exists when + // the processor runs, so nothing to compact. + drain(t, r.Run(t.Context(), userID, sessionID, genai.NewContentFromText("q1", genai.RoleUser), agent.RunConfig{})) + if got := summarizer.calls(); got != 0 { + t.Fatalf("summarizer ran %d times on the first turn, want 0", got) + } + + // Second turn: history now holds q1/answer 1/q2 plus a reported token + // count, so the processor compacts before the model call. + drain(t, r.Run(t.Context(), userID, sessionID, genai.NewContentFromText("q2", genai.RoleUser), agent.RunConfig{})) + if got := summarizer.calls(); got == 0 { + t.Fatal("summarizer never ran on the second turn, want tail-retention compaction") + } + + if got := len(compactionEventsIn(getSession(t, svc, userID, sessionID))); got == 0 { + t.Fatal("no compaction event was persisted") + } + + // The compaction landed before contents were built, so this very turn's + // prompt already carries the summary instead of the compacted turn. + prompt := promptText(m.lastPrompt()) + if !strings.Contains(prompt, "TAIL-SUMMARY") { + t.Errorf("prompt does not contain the summary:\n%s", prompt) + } + if strings.Contains(prompt, "q1") { + t.Errorf("prompt still contains the compacted turn q1:\n%s", prompt) + } +} + +func TestRunnerTailRetentionRespectsThreshold(t *testing.T) { + t.Parallel() + + const userID, sessionID = "u", "s" + // Reported prompts stay well under the threshold, so nothing compacts no + // matter how many turns accumulate. + m := &usageModel{promptTokens: 10} + summarizer := &recordingSummarizer{summary: "unused"} + r, svc := newCompactionRunner(t, m, &compaction.Config{ + TokenThreshold: 1000, + EventRetentionSize: 1, + Summarizer: summarizer, + }) + + for range 5 { + drain(t, r.Run(t.Context(), userID, sessionID, genai.NewContentFromText("q", genai.RoleUser), agent.RunConfig{})) + } + + if got := summarizer.calls(); got != 0 { + t.Errorf("summarizer ran %d times below the token threshold, want 0", got) + } + if got := len(compactionEventsIn(getSession(t, svc, userID, sessionID))); got != 0 { + t.Errorf("session holds %d compaction events below the threshold, want 0", got) + } +} + +// TestRunnerTailRetentionFailureDoesNotAbortTheTurn checks that a mid-turn +// compaction failure degrades to a larger prompt rather than killing the turn. +// +// Tail retention runs before a model call, inside an invocation whose tools may +// already have run and committed their side effects. Aborting there costs the +// user an answer, leaves the side effects standing, and orphans any summary +// already written, all to report that an optimisation did not happen. The +// threshold sits well below the real context limit, so the call usually still +// succeeds; when it does not, the provider's own error says more. +func TestRunnerTailRetentionFailureDoesNotAbortTheTurn(t *testing.T) { + t.Parallel() + + const userID, sessionID = "u", "s" + m := &usageModel{promptTokens: 5000} + r, svc := newCompactionRunner(t, m, &compaction.Config{ + TokenThreshold: 1000, + EventRetentionSize: 1, + Summarizer: failingSummarizer{}, + }) + + drain(t, r.Run(t.Context(), userID, sessionID, genai.NewContentFromText("q1", genai.RoleUser), agent.RunConfig{})) + + // The second turn trips the threshold and the summarizer fails. + var gotErr error + events := 0 + for _, err := range r.Run(t.Context(), userID, sessionID, genai.NewContentFromText("q2", genai.RoleUser), agent.RunConfig{}) { + if err != nil { + gotErr = err + break + } + events++ + } + if gotErr != nil { + t.Errorf("a failed mid-turn compaction aborted the turn: %v", gotErr) + } + if events == 0 { + t.Error("the turn produced no events, so the user got no answer") + } + // Nothing was recorded, so the next turn tries again rather than believing + // history was compacted. + if got := len(compactionEventsIn(getSession(t, svc, userID, sessionID))); got != 0 { + t.Errorf("stored %d compaction events despite the summarizer failing", got) + } +} + +func TestRunnerBothStrategiesCoexist(t *testing.T) { + t.Parallel() + + const userID, sessionID = "u", "s" + // Both triggers are armed: tail retention fires mid-turn on the reported + // token count, sliding window fires after every completed turn. A turn + // compacted mid-flight is not compacted again when it ends, so this + // exercises the hand-off between the two. + m := &usageModel{promptTokens: 5000} + summarizer := &recordingSummarizer{summary: "SUMMARY"} + r, svc := newCompactionRunner(t, m, &compaction.Config{ + TokenThreshold: 1000, + EventRetentionSize: 2, + CompactionInterval: 1, + Summarizer: summarizer, + }) + + for i := range 4 { + drain(t, r.Run(t.Context(), userID, sessionID, + genai.NewContentFromText(fmt.Sprintf("q%d", i), genai.RoleUser), agent.RunConfig{})) + } + + sess := getSession(t, svc, userID, sessionID) + if got := len(compactionEventsIn(sess)); got == 0 { + t.Fatal("no compaction events were produced with both strategies enabled") + } + + // Whatever mix of summaries accumulated, the prompt must stay coherent: + // every surviving compaction range is honoured and nothing is duplicated. + var events []*session.Event + for ev := range sess.Events().All() { + events = append(events, ev) + } + applied := compactioninternal.Apply(events) + + seen := make(map[string]bool) + for _, ev := range applied { + if ev.ID == "" { + continue + } + if seen[ev.ID] { + t.Errorf("event %q appears twice in the compacted prompt", ev.ID) + } + seen[ev.ID] = true + } + if len(applied) >= len(events) { + t.Errorf("compaction did not shrink history: %d events in, %d out", len(events), len(applied)) + } + + // The newest summary must not be subsumed, or the prompt would lose it. + if latest := compactioninternal.LatestCompactionEvent(events); latest == nil { + t.Error("no surviving compaction event; every summary was subsumed") + } +} + +// cancelingSummarizer fails with an error that wraps context.Canceled, which is +// what a summarizer whose own context died looks like. +type cancelingSummarizer struct{} + +func (s *cancelingSummarizer) SummarizeEvents(_ context.Context, _ []*session.Event) (compaction.SummarizeResult, error) { + return compaction.SummarizeResult{}, fmt.Errorf("summarizer model call failed: %w", context.Canceled) +} + +// TestTailRetentionCancelledSummarizerLeavesTheTurnIntact covers the case that +// used to produce the worst possible outcome. +// +// A summarizer failing on a cancelled context yielded an error whose chain +// contained context.Canceled, and the workflow scheduler drops those, so the +// turn ended with no answer, no events and no error either. Mid-turn compaction +// failures are no longer yielded at all, so the turn simply runs on. +func TestTailRetentionCancelledSummarizerLeavesTheTurnIntact(t *testing.T) { + t.Parallel() + + const userID, sessionID = "u", "s" + + m := &usageModel{promptTokens: 5000} + r, _ := newCompactionRunner(t, m, &compaction.Config{ + TokenThreshold: 100, + EventRetentionSize: 1, + Summarizer: &cancelingSummarizer{}, + }) + + drain(t, r.Run(t.Context(), userID, sessionID, genai.NewContentFromText("q1", genai.RoleUser), agent.RunConfig{})) + + var gotErr error + events := 0 + for _, err := range r.Run(t.Context(), userID, sessionID, genai.NewContentFromText("q2", genai.RoleUser), agent.RunConfig{}) { + if err != nil { + gotErr = err + break + } + events++ + } + if gotErr != nil { + t.Errorf("unexpected error from the turn: %v", gotErr) + } + if events == 0 { + t.Fatal("the turn produced no events and no error, which is the empty-turn outcome this guards against") + } +} + +// TestTailRetentionStandsDownTheSlidingWindow checks that a turn compacted +// mid-flight is not summarized a second time the moment it ends. +// +// The two strategies are independent triggers on the same history. Without a +// hand-off, a turn that crossed the token threshold pays for a second model +// call to re-summarize what was just summarized, and leaves two ranges over +// overlapping spans. The reference implementation avoids this by evaluating +// both in one place and returning early; here the mid-turn pass records that it +// ran and the post-invocation pass stands down. +func TestTailRetentionStandsDownTheSlidingWindow(t *testing.T) { + t.Parallel() + + const userID, sessionID = "u", "s" + + // Tuned so both strategies want to fire on the same turn, which is the only + // arrangement that exercises the hand-off. Interval 2 keeps the sliding + // window quiet on turn 1 so history can accumulate; by turn 2 there are + // three events, which is more than the retained tail, so tail retention + // fires mid-turn, and two completed invocations, so the sliding window + // would fire the moment the turn ends. + m := &usageModel{promptTokens: 5000} + summarizer := &recordingSummarizer{summary: "SUMMARY"} + r, svc := newCompactionRunner(t, m, &compaction.Config{ + TokenThreshold: 1000, + EventRetentionSize: 2, + CompactionInterval: 2, + Summarizer: summarizer, + }) + + perTurn := make([]int, 0, 2) + prev := 0 + for i := range 2 { + drain(t, r.Run(t.Context(), userID, sessionID, + genai.NewContentFromText(fmt.Sprintf("q%d", i), genai.RoleUser), agent.RunConfig{})) + perTurn = append(perTurn, summarizer.calls()-prev) + prev = summarizer.calls() + } + + // Turn 1 compacts nothing. Turn 2 compacts exactly once: tail retention + // mid-flight, and then the sliding window stands down. + if perTurn[0] != 0 || perTurn[1] != 1 { + t.Errorf("summarizer calls per turn = %v, want [0 1]: the second turn was compacted twice", perTurn) + } + if got := len(compactionEventsIn(getSession(t, svc, userID, sessionID))); got != 1 { + t.Errorf("stored %d compaction events, want 1", got) + } +} + +// toolLoopModel calls a tool repeatedly, then answers, always reporting the +// same prompt size. It stands in for a long tool loop whose retained tail alone +// already exceeds the threshold, so compacting cannot bring the prompt down. +type toolLoopModel struct { + mu sync.Mutex + calls int + rounds int + tokens int32 +} + +func (m *toolLoopModel) Name() string { return "tool-loop" } + +func (m *toolLoopModel) GenerateContent(_ context.Context, _ *model.LLMRequest, _ bool) iter.Seq2[*model.LLMResponse, error] { + m.mu.Lock() + m.calls++ + n := m.calls + m.mu.Unlock() + + return func(yield func(*model.LLMResponse, error) bool) { + usage := &genai.GenerateContentResponseUsageMetadata{PromptTokenCount: m.tokens} + if n <= m.rounds { + yield(&model.LLMResponse{ + Content: &genai.Content{Role: "model", Parts: []*genai.Part{ + {FunctionCall: &genai.FunctionCall{ID: fmt.Sprintf("c%d", n), Name: "ping"}}, + }}, + UsageMetadata: usage, + }, nil) + return + } + yield(&model.LLMResponse{Content: genai.NewContentFromText("done", "model"), UsageMetadata: usage}, nil) + } +} + +// TestTailRetentionStopsWhenItIsNotHelping checks that compaction gives up +// inside a turn once it stops reducing the prompt. +// +// The threshold is crossed before every model call in a tool loop. If the +// retained tail alone already exceeds it, compacting summarizes a little more +// each round and leaves the prompt exactly as far over, so every round pays for +// a summarizer call that changes nothing. +func TestTailRetentionStopsWhenItIsNotHelping(t *testing.T) { + t.Parallel() + + const userID, sessionID = "u", "s" + + ping, err := functiontool.New(functiontool.Config{Name: "ping", Description: "returns pong"}, + func(_ agent.Context, _ struct{}) (string, error) { return "pong", nil }) + if err != nil { + t.Fatalf("functiontool.New() error = %v", err) + } + + m := &toolLoopModel{rounds: 6, tokens: 5000} + root, err := llmagent.New(llmagent.Config{Name: "assistant", Model: m, Tools: []tool.Tool{ping}}) + if err != nil { + t.Fatalf("llmagent.New() error = %v", err) + } + summarizer := &recordingSummarizer{summary: "SUMMARY"} + r, err := New(Config{ + AppName: "compaction_app", + Agent: root, + SessionService: session.InMemoryService(), + AutoCreateSession: true, + Compaction: &compaction.Config{ + TokenThreshold: 1000, + EventRetentionSize: 2, + Summarizer: summarizer, + }, + }) + if err != nil { + t.Fatalf("New() error = %v", err) + } + + drain(t, r.Run(t.Context(), userID, sessionID, genai.NewContentFromText("go", genai.RoleUser), agent.RunConfig{})) + + // One attempt is right: it is worth trying once. Repeating is not, because + // the reported prompt never falls. + if got := summarizer.calls(); got > 1 { + t.Errorf("summarizer ran %d times in one turn while the prompt never shrank, want at most 1", got) + } + if m.calls < 3 { + t.Fatalf("the model only ran %d times, so the tool loop did not happen and this proved nothing", m.calls) + } +} + +// TestRunnerDefaultSummarizerIsBounded pins that the summarizer the runner +// installs cannot hold a turn open indefinitely. +// +// Compaction runs inside the run loop, and the post-invocation pass runs from a +// defer, so a provider that never answers parks the turn behind it. The +// Timeout field's own documentation says it is worth setting, and the one +// summarizer an application did not configure was the one without it. func TestRunnerDefaultSummarizerIsBounded(t *testing.T) { const userID, sessionID = "u", "s" @@ -1090,6 +1479,120 @@ func (m *hangingSummarizerModel) GenerateContent(ctx context.Context, _ *model.L // // A bound is the whole claim the package documentation makes for this strategy, // so it is asserted directly rather than inferred from a compaction happening. +func TestTailRetentionKeepsThePromptBounded(t *testing.T) { + t.Parallel() + + const userID, sessionID = "u", "s" + // Reports a count derived from the prompt it was given, the way a real + // model does. A fixed count would make the progress gate correctly conclude + // that compaction never helps and latch off, which is a different test. + m := &proportionalUsageModel{} + // A summary the size a real one is, roughly 500 characters, rather than a + // short marker. Size is what makes this test load-bearing: the failure it + // guards against is one summary per pass surviving into the prompt instead + // of each superseding the last, and with a seven-character summary sixty + // turns of that is still a small prompt, so the assertion below passes + // while the property is broken. At this length the same defect measured + // 24,991 characters against 551. + summaryText := strings.Repeat("summary text ", 40) + r, _ := newCompactionRunner(t, m, &compaction.Config{ + TokenThreshold: 200, + EventRetentionSize: 2, + Summarizer: &recordingSummarizer{summary: summaryText}, + }) + + var early, late int + const rounds = 60 + for i := range rounds { + drain(t, r.Run(t.Context(), userID, sessionID, + genai.NewContentFromText(fmt.Sprintf("question %d", i), genai.RoleUser), agent.RunConfig{})) + + size := 0 + for _, c := range m.lastPrompt() { + for _, p := range c.Parts { + size += len(p.Text) + } + } + switch i { + case rounds / 3: + early = size + case rounds - 1: + late = size + } + } + + // Some slack, because a rolling summary and its raw tail vary in length + // from turn to turn. What must not happen is growth proportional to the + // number of turns. + if late > early*2 { + t.Errorf("prompt grew from %d characters at turn %d to %d at turn %d: tail retention is not bounding it", + early, rounds/3, late, rounds-1) + } + + // The mechanism, stated separately from the symptom. A rolling summary is + // supposed to replace the one it was built from, so however many passes + // ran, one summary reaches the model. Counting them says which way a size + // regression went, and catches the accumulation before it is large enough + // to move the total. + summaries := 0 + for _, c := range m.lastPrompt() { + for _, p := range c.Parts { + if strings.Contains(p.Text, summaryText) { + summaries++ + } + } + } + if summaries > 1 { + t.Errorf("final prompt carries %d summaries, want 1: each pass is adding a summary rather than superseding the last", summaries) + } +} + +// proportionalUsageModel reports a prompt token count derived from the prompt it +// received, so compaction visibly shrinks the next reading. +type proportionalUsageModel struct { + mu sync.Mutex + prompts [][]*genai.Content +} + +func (m *proportionalUsageModel) Name() string { return "proportional" } + +func (m *proportionalUsageModel) GenerateContent(_ context.Context, req *model.LLMRequest, _ bool) iter.Seq2[*model.LLMResponse, error] { + m.mu.Lock() + m.prompts = append(m.prompts, req.Contents) + n := len(m.prompts) + m.mu.Unlock() + + chars := 0 + for _, c := range req.Contents { + for _, p := range c.Parts { + chars += len(p.Text) + } + } + return func(yield func(*model.LLMResponse, error) bool) { + yield(&model.LLMResponse{ + Content: genai.NewContentFromText(fmt.Sprintf("answer %d", n), "model"), + UsageMetadata: &genai.GenerateContentResponseUsageMetadata{PromptTokenCount: int32(chars)}, + }, nil) + } +} + +func (m *proportionalUsageModel) lastPrompt() []*genai.Content { + m.mu.Lock() + defer m.mu.Unlock() + if len(m.prompts) == 0 { + return nil + } + return m.prompts[len(m.prompts)-1] +} + +// TestPluginCannotSmuggleAFunctionCallIntoASummary pins that a plugin's +// replacement summary is filtered like a summarizer's. +// +// A plugin may see and rewrite a summary before it is stored, which is the +// point of routing it through the pipeline. Its replacement went to the session +// unexamined, so content carrying a text part and a FunctionCall reached a real +// model prompt as an unpaired call, which is exactly what the filter on the +// summarizer path exists to stop. func TestPluginCannotSmuggleAFunctionCallIntoASummary(t *testing.T) { t.Parallel()