From 587d2525c307bedebe847e01ad463a353881bb7d Mon Sep 17 00:00:00 2001 From: Google Team Member Date: Wed, 4 Nov 2026 10:18:15 -0800 Subject: [PATCH] fix: count events after the last model response toward the compaction threshold PiperOrigin-RevId: 1013953371 --- .../TailRetentionEventCompactor.java | 40 ++++-- .../adk/flows/llmflows/CompactionTest.java | 64 +++++++-- .../TailRetentionEventCompactorTest.java | 121 ++++++++++++++++++ 3 files changed, 202 insertions(+), 23 deletions(-) diff --git a/core/src/main/java/com/google/adk/summarizer/TailRetentionEventCompactor.java b/core/src/main/java/com/google/adk/summarizer/TailRetentionEventCompactor.java index 4159ad2cd..2945c9a85 100644 --- a/core/src/main/java/com/google/adk/summarizer/TailRetentionEventCompactor.java +++ b/core/src/main/java/com/google/adk/summarizer/TailRetentionEventCompactor.java @@ -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; @@ -41,7 +40,10 @@ *
  • Keeps the {@code retentionSize} most recent events raw. *
  • Compacts all events that never compacted and older than the retained tail, including the * most recent compaction event, into a new summary event. - *
  • Triggers compaction only if the prompt token count exceeds the {@code tokenThreshold}. + *
  • 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. *
  • The new summary event is generated by the {@link BaseEventSummarizer}. *
  • Appends this new summary event to the end of the event stream. * @@ -140,10 +142,11 @@ public Completable compact(Session session, BaseSessionService sessionService) { * @param events The list of events to process. */ private Maybe> getCompactionEvents(List events) { - Optional count = getLatestPromptTokenCount(events); + Optional 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(); @@ -226,13 +229,28 @@ private int estimateTokenCount(List events) { return events.stream().mapToInt(event -> event.stringifyContent().length()).sum() / 4; } - private Optional getLatestPromptTokenCount(List 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 getEstimatedPromptTokenCount(List events) { + ListIterator 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 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) { diff --git a/core/src/test/java/com/google/adk/flows/llmflows/CompactionTest.java b/core/src/test/java/com/google/adk/flows/llmflows/CompactionTest.java index 3ceba5641..a5017beb6 100644 --- a/core/src/test/java/com/google/adk/flows/llmflows/CompactionTest.java +++ b/core/src/test/java/com/google/adk/flows/llmflows/CompactionTest.java @@ -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; @@ -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(); @@ -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); @@ -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 @@ -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(); + } } diff --git a/core/src/test/java/com/google/adk/summarizer/TailRetentionEventCompactorTest.java b/core/src/test/java/com/google/adk/summarizer/TailRetentionEventCompactorTest.java index 7a4a3ddb9..c3d8840a4 100644 --- a/core/src/test/java/com/google/adk/summarizer/TailRetentionEventCompactorTest.java +++ b/core/src/test/java/com/google/adk/summarizer/TailRetentionEventCompactorTest.java @@ -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; @@ -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 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 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 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 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 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 events = @@ -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()