Skip to content

fix: count events after the last model response toward the compaction threshold #1717

New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Merged
copybara-service merged 1 commit into from
Nov 4, 2026
+202 −23
Merged

fix: count events after the last model response toward the compaction threshold #1717

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

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 29 additions & 11 deletions core/src/main/java/com/google/adk/summarizer/TailRetentionEventCompactor.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,6 @@
import com.google.adk.events.EventCompaction;
import com.google.adk.sessions.BaseSessionService;
import com.google.adk.sessions.Session;
import com.google.common.collect.Lists;
import com.google.genai.types.GenerateContentResponseUsageMetadata;
import io.reactivex.rxjava3.core.Completable;
import io.reactivex.rxjava3.core.Maybe;
Expand All @@ -41,7 +40,10 @@
* <li>Keeps the {@code retentionSize} most recent events raw.
* <li>Compacts all events that never compacted and older than the retained tail, including the
* most recent compaction event, into a new summary event.
* <li>Triggers compaction only if the prompt token count exceeds the {@code tokenThreshold}.
* <li>Triggers compaction only if the estimated prompt token count exceeds the {@code
* tokenThreshold}: the last reported prompt token count plus about one token per four
* characters of the response that reported it and the events after it. Without a count
* reported since the latest compaction, all events since then are estimated this way.
* <li>The new summary event is generated by the {@link BaseEventSummarizer}.
* <li>Appends this new summary event to the end of the event stream.
* </ul>
Expand Down Expand Up @@ -140,10 +142,11 @@ public Completable compact(Session session, BaseSessionService sessionService) {
* @param events The list of events to process.
*/
private Maybe<List<Event>> getCompactionEvents(List<Event> events) {
Optional<Integer> count = getLatestPromptTokenCount(events);
Optional<Integer> count = getEstimatedPromptTokenCount(events);
if (count.isPresent() && count.get() <= tokenThreshold) {
logger.debug(
"Skipping compaction. Prompt token count {} is within threshold {}",
"Skipping compaction. Prompt token count {} (last reported plus estimate) is within"
+ " threshold {}",
count.get(),
tokenThreshold);
return Maybe.empty();
Expand Down Expand Up @@ -226,13 +229,28 @@ private int estimateTokenCount(List<Event> events) {
return events.stream().mapToInt(event -> event.stringifyContent().length()).sum() / 4;
}

private Optional<Integer> getLatestPromptTokenCount(List<Event> events) {
return Lists.reverse(events).stream()
.map(Event::usageMetadata)
.flatMap(Optional::stream)
.map(GenerateContentResponseUsageMetadata::promptTokenCount)
.flatMap(Optional::stream)
.findFirst();
/**
* Returns the latest reported prompt token count plus an estimate for the events it does not
* cover, or empty if there is no such count after the latest compaction.
*/
private Optional<Integer> getEstimatedPromptTokenCount(List<Event> events) {
ListIterator<Event> iter = events.listIterator(events.size());
while (iter.hasPrevious()) {
Event event = iter.previous();
if (isCompactEvent(event)) {
// A count reported before a compaction describes a prompt that the compaction replaced.
return Optional.empty();
}
Optional<Integer> promptTokenCount =
event.usageMetadata().flatMap(GenerateContentResponseUsageMetadata::promptTokenCount);
if (promptTokenCount.isPresent()) {
// The count excludes this response and every event appended after it.
return Optional.of(
promptTokenCount.get()
+ estimateTokenCount(events.subList(iter.nextIndex(), events.size())));
}
}
return Optional.empty();
}

private static boolean isCompactEvent(Event event) {
Expand Down
64 changes: 52 additions & 12 deletions core/src/test/java/com/google/adk/flows/llmflows/CompactionTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,10 @@
import com.google.adk.summarizer.BaseEventSummarizer;
import com.google.adk.summarizer.EventsCompactionConfig;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.genai.types.Content;
import com.google.genai.types.GenerateContentResponseUsageMetadata;
import com.google.genai.types.Part;
import io.reactivex.rxjava3.core.Maybe;
import io.reactivex.rxjava3.core.Single;
import java.util.Optional;
Expand Down Expand Up @@ -82,13 +85,9 @@ public void processRequest_withConfig_triggersCompaction() {
when(context.eventsCompactionConfig()).thenReturn(Optional.of(config));

// Setup events with usage > 100 to trigger compaction
Event event1 = mock(Event.class);
Event event2 = mock(Event.class);
Event event3 = mock(Event.class);
when(event3.usageMetadata())
.thenReturn(
Optional.of(
GenerateContentResponseUsageMetadata.builder().promptTokenCount(200).build()));
Event event1 = Event.builder().timestamp(1).build();
Event event2 = Event.builder().timestamp(2).build();
Event event3 = eventWithPromptTokenCount(3, 200);

session =
Session.builder("test-session").events(ImmutableList.of(event1, event2, event3)).build();
Expand Down Expand Up @@ -117,11 +116,7 @@ public void processRequest_withConfig_skipsCompactionIfBelowThreshold() {
when(context.eventsCompactionConfig()).thenReturn(Optional.of(config));

// Setup events with usage 200 (below 500)
Event event3 = mock(Event.class);
when(event3.usageMetadata())
.thenReturn(
Optional.of(
GenerateContentResponseUsageMetadata.builder().promptTokenCount(200).build()));
Event event3 = eventWithPromptTokenCount(3, 200);

session = Session.builder("test-session").events(ImmutableList.of(event3)).build();
when(context.session()).thenReturn(session);
Expand All @@ -137,6 +132,41 @@ public void processRequest_withConfig_skipsCompactionIfBelowThreshold() {
verify(sessionService, never()).appendEvent(any(), any());
}

@Test
public void processRequest_withConfig_toolResponseAfterTokenUsage_triggersCompaction() {
// Threshold 300: the reported count of 200 alone is within it, but the function response
// appended after it (an 800-character result, over 200 tokens) pushes the estimate above.
EventsCompactionConfig config = new EventsCompactionConfig(5, 1, summarizer, 300, 2);
when(context.eventsCompactionConfig()).thenReturn(Optional.of(config));

Event event1 = Event.builder().timestamp(1).build();
Event event2 = eventWithPromptTokenCount(2, 200);
Event event3 =
Event.builder()
.timestamp(3)
.content(
Content.fromParts(
Part.fromFunctionResponse("tool", ImmutableMap.of("result", "x".repeat(800)))))
.build();

session =
Session.builder("test-session").events(ImmutableList.of(event1, event2, event3)).build();
when(context.session()).thenReturn(session);

Event summaryEvent = mock(Event.class);
when(summarizer.summarizeEvents(any())).thenReturn(Maybe.just(summaryEvent));
when(sessionService.appendEvent(any(), any())).thenReturn(Single.just(summaryEvent));

Compaction compaction = new Compaction();
compaction
.processRequest(context, request)
.test()
.assertNoErrors()
.assertValue(r -> r.updatedRequest() == request);

verify(sessionService).appendEvent(eq(session), eq(summaryEvent));
}

@Test
public void processRequest_withConfig_nullRetentionSize_doesNothing() {
// Setup config with retentionSize = null
Expand All @@ -154,4 +184,14 @@ public void processRequest_withConfig_nullRetentionSize_doesNothing() {
verify(sessionService, never()).appendEvent(any(), any());
verify(context, never()).session();
}

private static Event eventWithPromptTokenCount(long timestamp, int promptTokenCount) {
return Event.builder()
.timestamp(timestamp)
.usageMetadata(
GenerateContentResponseUsageMetadata.builder()
.promptTokenCount(promptTokenCount)
.build())
.build();
}
}
121 changes: 121 additions & 0 deletions core/src/test/java/com/google/adk/summarizer/TailRetentionEventCompactorTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@
import com.google.adk.sessions.BaseSessionService;
import com.google.adk.sessions.Session;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.genai.types.Content;
import com.google.genai.types.GenerateContentResponseUsageMetadata;
import com.google.genai.types.Part;
Expand Down Expand Up @@ -160,6 +161,118 @@ public void compaction_happensWhenTokenUsageAboveThreshold() {
verify(mockSessionService).appendEvent(eq(session), eq(summaryEvent));
}

@Test
public void compaction_happensWhenEventsAfterTokenUsagePushAboveThreshold() {
// Threshold is 300, the last model call used 200 prompt tokens, and a tool response of
// 800 characters (about 200 tokens) arrived after it: 200 + 200 > 300.
EventCompactor compactor = new TailRetentionEventCompactor(mockSummarizer, 1, 300);
ImmutableList<Event> events =
ImmutableList.of(
createEvent(1, "Event1"),
withUsage(createEvent(2, "Call"), 200),
createEvent(3, "x".repeat(800)));
Session session = Session.builder("id").events(events).build();
Event summaryEvent = createEvent(4, "Summary");

when(mockSummarizer.summarizeEvents(any())).thenReturn(Maybe.just(summaryEvent));
when(mockSessionService.appendEvent(any(), any())).thenReturn(Single.just(summaryEvent));

compactor.compact(session, mockSessionService).blockingSubscribe();

verify(mockSummarizer).summarizeEvents(eventListCaptor.capture());
assertThat(
eventListCaptor.getValue().stream().map(TailRetentionEventCompactorTest::getPromptText))
.containsExactly("Event1", "Call")
.inOrder();
verify(mockSessionService).appendEvent(eq(session), eq(summaryEvent));
}

@Test
public void compaction_happensWhenToolResponseAfterTokenUsagePushesAboveThreshold() {
// Threshold is 300, the last model call used 200 prompt tokens, and a function response with
// an 800-character result (over 200 tokens once stringified) arrived after it.
EventCompactor compactor = new TailRetentionEventCompactor(mockSummarizer, 1, 300);
ImmutableList<Event> events =
ImmutableList.of(
createEvent(1, "Event1"),
withUsage(createEvent(2, "Call"), 200),
createToolResponseEvent(3, "x".repeat(800)));
Session session = Session.builder("id").events(events).build();
Event summaryEvent = createEvent(4, "Summary");

when(mockSummarizer.summarizeEvents(any())).thenReturn(Maybe.just(summaryEvent));
when(mockSessionService.appendEvent(any(), any())).thenReturn(Single.just(summaryEvent));

compactor.compact(session, mockSessionService).blockingSubscribe();

verify(mockSummarizer).summarizeEvents(eventListCaptor.capture());
assertThat(
eventListCaptor.getValue().stream().map(TailRetentionEventCompactorTest::getPromptText))
.containsExactly("Event1", "Call")
.inOrder();
verify(mockSessionService).appendEvent(eq(session), eq(summaryEvent));
}

@Test
public void compaction_skippedWhenEventsAfterTokenUsageStayBelowThreshold() {
// Threshold is 300, the last model call used 200 prompt tokens, and a tool response of
// 200 characters (about 50 tokens) arrived after it: 200 + 50 <= 300.
EventCompactor compactor = new TailRetentionEventCompactor(mockSummarizer, 1, 300);
ImmutableList<Event> events =
ImmutableList.of(
createEvent(1, "Event1"),
withUsage(createEvent(2, "Call"), 200),
createEvent(3, "x".repeat(200)));
Session session = Session.builder("id").events(events).build();

compactor.compact(session, mockSessionService).blockingSubscribe();

verify(mockSummarizer, never()).summarizeEvents(any());
verify(mockSessionService, never()).appendEvent(any(), any());
}

@Test
public void compaction_countsTheResponseThatReportedTokenUsage() {
// promptTokenCount excludes the response that reports it. Its 40 characters (~10 tokens) are
// what push 200 + 100 (the 400-character tool response) past the threshold of 300.
EventCompactor compactor = new TailRetentionEventCompactor(mockSummarizer, 1, 300);
ImmutableList<Event> events =
ImmutableList.of(
createEvent(1, "Event1"),
withUsage(createEvent(2, "y".repeat(40)), 200),
createEvent(3, "x".repeat(400)));
Session session = Session.builder("id").events(events).build();
Event summaryEvent = createEvent(4, "Summary");

when(mockSummarizer.summarizeEvents(any())).thenReturn(Maybe.just(summaryEvent));
when(mockSessionService.appendEvent(any(), any())).thenReturn(Single.just(summaryEvent));

compactor.compact(session, mockSessionService).blockingSubscribe();

verify(mockSummarizer).summarizeEvents(any());
}

@Test
public void compaction_ignoresTokenUsageReportedBeforeLatestCompaction() {
// The 1000-token count describes the prompt that the compaction replaced. What remains
// ("Summary", "Retain1", "Retain2") is ~5 estimated tokens, within the threshold of 300.
EventCompactor compactor = new TailRetentionEventCompactor(mockSummarizer, 1, 300);
ImmutableList<Event> events =
ImmutableList.of(
createEvent(1, "Event1"),
withUsage(createEvent(2, "Call"), 1000),
createEvent(3, "Retain1"),
createCompactedEvent(
/* startTimestamp= */ 1, /* endTimestamp= */ 2, "Summary", /* eventTimestamp= */ 4),
createEvent(5, "Retain2"));
Session session = Session.builder("id").events(events).build();

compactor.compact(session, mockSessionService).blockingSubscribe();

verify(mockSummarizer, never()).summarizeEvents(any());
verify(mockSessionService, never()).appendEvent(any(), any());
}

@Test
public void compact_notEnoughEvents_doesNothing() {
ImmutableList<Event> events =
Expand Down Expand Up @@ -330,6 +443,14 @@ private static Event createEvent(long timestamp, String text) {
.build();
}

private static Event createToolResponseEvent(long timestamp, String result) {
return Event.builder()
.timestamp(timestamp)
.content(
Content.fromParts(Part.fromFunctionResponse("tool", ImmutableMap.of("result", result))))
.build();
}

private static String getPromptText(Event event) {
return event
.content()
Expand Down
Loading