From fe96d9761fcce62de01678013c7ae43b6c219eba Mon Sep 17 00:00:00 2001 From: Wiktoria Walczak Date: Mon, 22 Jun 2026 04:56:22 -0700 Subject: [PATCH] fix: exclude rewound invocations from sliding-window compaction PiperOrigin-RevId: 935992227 --- .../com/google/adk/kt/events/RewindEvents.kt | 49 ++++++++++++++++ .../kt/processors/HistoryRewriterProcessor.kt | 29 +--------- .../summarizer/SlidingWindowEventCompactor.kt | 7 ++- .../SlidingWindowEventCompactorTest.kt | 58 +++++++++++++++++++ .../com/google/adk/kt/testing/TestEvent.kt | 9 +++ 5 files changed, 123 insertions(+), 29 deletions(-) create mode 100644 core/src/commonMain/kotlin/com/google/adk/kt/events/RewindEvents.kt diff --git a/core/src/commonMain/kotlin/com/google/adk/kt/events/RewindEvents.kt b/core/src/commonMain/kotlin/com/google/adk/kt/events/RewindEvents.kt new file mode 100644 index 00000000..8c0a554a --- /dev/null +++ b/core/src/commonMain/kotlin/com/google/adk/kt/events/RewindEvents.kt @@ -0,0 +1,49 @@ +/* + * 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 com.google.adk.kt.events + +/** + * Returns [events] with rewound invocations removed. + * + * Iterates backward. When an event carries `actions.rewindBeforeInvocationId == X`, drops that + * event together with every event between it and the earliest event of invocation `X` (inclusive), + * then resumes the backward walk from there. + * + * This is the single source of truth for "which events are live" after rewinds. Both LLM prompt + * building ([com.google.adk.kt.processors.HistoryRewriterProcessor]) and context compaction must + * agree on it, otherwise rewound content can leak back into prompts through a compaction summary. + */ +internal fun applyRewinds(events: List): List { + val kept = mutableListOf() + var i = events.size - 1 + while (i >= 0) { + val event = events[i] + val rewindInvocationId = event.actions.rewindBeforeInvocationId + if (!rewindInvocationId.isNullOrEmpty()) { + for (j in 0 until i) { + if (events[j].invocationId == rewindInvocationId) { + i = j + break + } + } + } else { + kept.add(event) + } + i-- + } + return kept.asReversed() +} diff --git a/core/src/commonMain/kotlin/com/google/adk/kt/processors/HistoryRewriterProcessor.kt b/core/src/commonMain/kotlin/com/google/adk/kt/processors/HistoryRewriterProcessor.kt index a5ac0e64..d5699d1a 100644 --- a/core/src/commonMain/kotlin/com/google/adk/kt/processors/HistoryRewriterProcessor.kt +++ b/core/src/commonMain/kotlin/com/google/adk/kt/processors/HistoryRewriterProcessor.kt @@ -17,6 +17,7 @@ package com.google.adk.kt.processors import com.google.adk.kt.agents.LlmAgent.IncludeContents import com.google.adk.kt.events.Event +import com.google.adk.kt.events.applyRewinds import com.google.adk.kt.serialization.Json import com.google.adk.kt.types.Content import com.google.adk.kt.types.FunctionCall @@ -92,34 +93,6 @@ internal class HistoryRewriterProcessor { return emptyList() } - /** - * Returns [events] with rewound invocations removed. - * - * Iterates backward. When an event carries `actions.rewindBeforeInvocationId == X`, drops that - * event together with every event between it and the earliest event of invocation `X` - * (inclusive), then resumes the backward walk from there. - */ - private fun applyRewinds(events: List): List { - val kept = mutableListOf() - var i = events.size - 1 - while (i >= 0) { - val event = events[i] - val rewindInvocationId = event.actions.rewindBeforeInvocationId - if (!rewindInvocationId.isNullOrEmpty()) { - for (j in 0 until i) { - if (events[j].invocationId == rewindInvocationId) { - i = j - break - } - } - } else { - kept.add(event) - } - i-- - } - return kept.asReversed() - } - /** * Returns whether [event] qualifies as the start of the current turn for [agentName] on * [currentBranch]: it must be visible in this agent's context, and it must be a user input or diff --git a/core/src/commonMain/kotlin/com/google/adk/kt/summarizer/SlidingWindowEventCompactor.kt b/core/src/commonMain/kotlin/com/google/adk/kt/summarizer/SlidingWindowEventCompactor.kt index 9457a25c..7f20500f 100644 --- a/core/src/commonMain/kotlin/com/google/adk/kt/summarizer/SlidingWindowEventCompactor.kt +++ b/core/src/commonMain/kotlin/com/google/adk/kt/summarizer/SlidingWindowEventCompactor.kt @@ -16,6 +16,7 @@ package com.google.adk.kt.summarizer import com.google.adk.kt.events.Event +import com.google.adk.kt.events.applyRewinds import com.google.adk.kt.logging.LoggerFactory import com.google.adk.kt.sessions.Session import com.google.adk.kt.sessions.SessionService @@ -48,7 +49,11 @@ class SlidingWindowEventCompactor(private val config: EventsCompactionConfig) : if (!config.hasSlidingWindowConfig()) return val summarizer = requireNotNull(config.summarizer) { "Missing EventSummarizer for event compaction." } - val compactionWindow = selectCompactionWindow(session.events) ?: return + // Drop rewound invocations first so the summary covers only live events. This keeps the + // compactor consistent with prompt building (HistoryRewriterProcessor also applies rewinds); + // otherwise rewound content would leak back into future prompts via the compaction summary. + val liveEvents = applyRewinds(session.events) + val compactionWindow = selectCompactionWindow(liveEvents) ?: return val compactionEvent = summarizer.summarizeEvents(compactionWindow) ?: return val appendedEvent = sessionService.appendEvent(session, compactionEvent) logger.debug { diff --git a/core/src/commonTest/kotlin/com/google/adk/kt/summarizer/SlidingWindowEventCompactorTest.kt b/core/src/commonTest/kotlin/com/google/adk/kt/summarizer/SlidingWindowEventCompactorTest.kt index c64037a8..bb1ccbb1 100644 --- a/core/src/commonTest/kotlin/com/google/adk/kt/summarizer/SlidingWindowEventCompactorTest.kt +++ b/core/src/commonTest/kotlin/com/google/adk/kt/summarizer/SlidingWindowEventCompactorTest.kt @@ -27,6 +27,7 @@ import com.google.adk.kt.testing.eventWithFunctionCall import com.google.adk.kt.testing.eventWithFunctionResponse import com.google.adk.kt.testing.eventWithHitlRequest import com.google.adk.kt.testing.modelEvent +import com.google.adk.kt.testing.rewindEvent import com.google.adk.kt.testing.testSession import com.google.adk.kt.testing.userEvent import kotlin.test.Test @@ -397,6 +398,63 @@ class SlidingWindowEventCompactorTest { assertEquals(listOf(firstUser, firstModel), summarizer.calls.single()) } + @Test + fun compact_rewoundInvocation_excludedFromSummary() = runTest { + val summarizer = RecordingSummarizer(returning = compactionEvent(startTs = 100L, endTs = 310L)) + val sessionService = RecordingSessionService() + val compactor = + SlidingWindowEventCompactor( + EventsCompactionConfig(compactionInterval = 2, overlapSize = 1, summarizer = summarizer) + ) + val session = testSession() + val firstUser = userEvent("user event to keep", invocationId = "inv_1", timestamp = 100L) + val firstModel = modelEvent("model response", invocationId = "inv_1", timestamp = 110L) + // inv_2 was rewound by the user; its content must never reach the summarizer. + val rewoundUser = userEvent("REWOUND_EVENT", invocationId = "inv_2", timestamp = 200L) + val rewoundModel = modelEvent("rewound reply", invocationId = "inv_2", timestamp = 210L) + val rewindMarker = + rewindEvent(invocationId = "rewind_inv", rewoundInvocationId = "inv_2", timestamp = 220L) + val thirdUser = userEvent("user event to keep", invocationId = "inv_3", timestamp = 300L) + val thirdModel = modelEvent("model response", invocationId = "inv_3", timestamp = 310L) + session.events.addAll( + listOf(firstUser, firstModel, rewoundUser, rewoundModel, rewindMarker, thirdUser, thirdModel) + ) + + compactor.compact(session, sessionService) + + // Only the two live invocations (inv_1, inv_3) are summarized; the rewound invocation and the + // rewind marker are dropped before window selection, so the rewound content never leaks. + assertEquals(listOf(firstUser, firstModel, thirdUser, thirdModel), summarizer.calls.single()) + // Raw events stay persisted; compaction reads through rewinds but does not delete history. + assertTrue(session.events.containsAll(listOf(rewoundUser, rewoundModel, rewindMarker))) + } + + @Test + fun compact_rewoundInvocationDoesNotCountTowardThreshold() = runTest { + val summarizer = RecordingSummarizer() + val sessionService = RecordingSessionService() + val compactor = + SlidingWindowEventCompactor( + EventsCompactionConfig(compactionInterval = 2, overlapSize = 1, summarizer = summarizer) + ) + val session = testSession() + // One live invocation plus one rewound invocation: only 1 live invocation < interval (2). + session.events.add(userEvent("live", invocationId = "inv_1", timestamp = 100L)) + session.events.add(modelEvent("ok", invocationId = "inv_1", timestamp = 110L)) + session.events.add(userEvent("REWOUND_SECRET", invocationId = "inv_2", timestamp = 200L)) + session.events.add(modelEvent("rewound reply", invocationId = "inv_2", timestamp = 210L)) + session.events.add( + rewindEvent(invocationId = "rewind_inv", rewoundInvocationId = "inv_2", timestamp = 220L) + ) + + compactor.compact(session, sessionService) + + // Rewound invocations (and the marker) do not count toward the interval, so the threshold is + // not met and nothing is summarized. + assertTrue(summarizer.calls.isEmpty()) + assertTrue(sessionService.appended.isEmpty()) + } + @Test fun compact_nullSummarizer_throwsIllegalArgumentException() = runTest { val compactor = diff --git a/core/src/commonTest/kotlin/com/google/adk/kt/testing/TestEvent.kt b/core/src/commonTest/kotlin/com/google/adk/kt/testing/TestEvent.kt index 85408c11..876a29b9 100644 --- a/core/src/commonTest/kotlin/com/google/adk/kt/testing/TestEvent.kt +++ b/core/src/commonTest/kotlin/com/google/adk/kt/testing/TestEvent.kt @@ -86,6 +86,15 @@ fun eventWithHitlRequest(invocationId: String, timestamp: Long, callId: String): timestamp = timestamp, ) +/** A `user`-authored marker [Event] that rewinds history before [rewoundInvocationId]. */ +fun rewindEvent(invocationId: String, rewoundInvocationId: String, timestamp: Long = 0L): Event = + Event( + author = Role.USER, + invocationId = invocationId, + actions = EventActions(rewindBeforeInvocationId = rewoundInvocationId), + timestamp = timestamp, + ) + /** An [Event] carrying an [EventCompaction] [summary] spanning [startTs]..[endTs]. */ fun compactionEvent( startTs: Long,