From 8ccda14e476c3d84d0a3248af5b5547405bce51d Mon Sep 17 00:00:00 2001 From: Wiktoria Walczak Date: Mon, 20 Jul 2026 03:30:03 -0700 Subject: [PATCH] feat: Add ContextCacheRequestProcessor and register it in the LLM flow PiperOrigin-RevId: 950726013 --- .../com/google/adk/kt/agents/LlmAgent.kt | 2 + .../ContextCacheRequestProcessor.kt | 113 ++++++++++++ .../ContextCacheRequestProcessorTest.kt | 165 ++++++++++++++++++ 3 files changed, 280 insertions(+) create mode 100644 core/src/commonMain/kotlin/com/google/adk/kt/processors/ContextCacheRequestProcessor.kt create mode 100644 core/src/commonTest/kotlin/com/google/adk/kt/processors/ContextCacheRequestProcessorTest.kt diff --git a/core/src/commonMain/kotlin/com/google/adk/kt/agents/LlmAgent.kt b/core/src/commonMain/kotlin/com/google/adk/kt/agents/LlmAgent.kt index ca472c6e..f4ec2b7e 100644 --- a/core/src/commonMain/kotlin/com/google/adk/kt/agents/LlmAgent.kt +++ b/core/src/commonMain/kotlin/com/google/adk/kt/agents/LlmAgent.kt @@ -32,6 +32,7 @@ import com.google.adk.kt.processors.AgentTransferProcessor import com.google.adk.kt.processors.BasicRequestProcessor import com.google.adk.kt.processors.CompactionRequestProcessor import com.google.adk.kt.processors.ContentsProcessor +import com.google.adk.kt.processors.ContextCacheRequestProcessor import com.google.adk.kt.processors.InstructionsProcessor import com.google.adk.kt.processors.LlmRequestProcessor import com.google.adk.kt.processors.LlmResponseProcessor @@ -193,6 +194,7 @@ class LlmAgent( // request context. CompactionRequestProcessor(), ContentsProcessor(), + ContextCacheRequestProcessor(), AgentTransferProcessor(), OutputSchemaProcessor(), ) diff --git a/core/src/commonMain/kotlin/com/google/adk/kt/processors/ContextCacheRequestProcessor.kt b/core/src/commonMain/kotlin/com/google/adk/kt/processors/ContextCacheRequestProcessor.kt new file mode 100644 index 00000000..295a0cb0 --- /dev/null +++ b/core/src/commonMain/kotlin/com/google/adk/kt/processors/ContextCacheRequestProcessor.kt @@ -0,0 +1,113 @@ +/* + * 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.processors + +import com.google.adk.kt.agents.InvocationContext +import com.google.adk.kt.events.Event +import com.google.adk.kt.logging.LoggerFactory +import com.google.adk.kt.models.CacheMetadata +import com.google.adk.kt.models.LlmRequest + +/** + * A request processor that enables context caching for the LLM request. + * + * When the invocation has a [InvocationContext.contextCacheConfig], this copies the config onto the + * request and recovers the most recent cache metadata and prompt token count from the current + * agent's session events, so the model-side cache manager (e.g. the Gemini cache manager) can + * validate and reuse an existing cache. When no config is present, the request is returned + * unchanged and caching stays disabled. The actual cache lifecycle is handled by the model. + */ +internal class ContextCacheRequestProcessor : LlmRequestProcessor { + override suspend fun process( + context: InvocationContext, + request: LlmRequest, + emitEvent: suspend (Event) -> Unit, + ): LlmRequest { + val cacheConfig = context.contextCacheConfig ?: return request + + val (cacheMetadata, previousTokenCount) = + findCacheInfoFromEvents(context.session.events, context.agent.name, context.invocationId) + + if (cacheMetadata != null) { + logger.debug { "Found cache metadata for agent ${context.agent.name}: $cacheMetadata" } + } + if (previousTokenCount != null) { + logger.debug { + "Found previous prompt token count for agent ${context.agent.name}: $previousTokenCount" + } + } + logger.debug { "Context caching enabled for agent ${context.agent.name}" } + + return request.copy( + cacheConfig = cacheConfig, + cacheMetadata = cacheMetadata, + cacheableContentsTokenCount = previousTokenCount, + ) + } + + /** + * Scans [events] (most recent first) for the latest cache metadata and prompt token count + * authored by [agentName]. + * + * The returned cache metadata has its [CacheMetadata.invocationsUsed] incremented by one when it + * comes from a different (earlier) invocation than [currentInvocationId] and refers to an active + * cache. This mirrors the Python ADK behavior of counting one additional reuse per invocation. + * + * @return A pair of the recovered cache metadata (or `null`) and the most recent prompt token + * count (or `null`). + */ + private fun findCacheInfoFromEvents( + events: List, + agentName: String, + currentInvocationId: String, + ): Pair { + var cacheMetadata: CacheMetadata? = null + var previousTokenCount: Int? = null + + for (event in events.asReversed()) { + if (event.author != agentName) continue + + if (cacheMetadata == null) { + val eventCacheMetadata = event.cacheMetadata + if (eventCacheMetadata != null) { + cacheMetadata = + if ( + event.invocationId != null && + event.invocationId != currentInvocationId && + eventCacheMetadata.cacheName != null + ) { + // Different invocation with an active cache: count one more reuse. + eventCacheMetadata.copy(invocationsUsed = eventCacheMetadata.invocationsUsed!! + 1) + } else { + eventCacheMetadata + } + } + } + + if (previousTokenCount == null) { + previousTokenCount = event.usageMetadata?.promptTokenCount + } + + if (cacheMetadata != null && previousTokenCount != null) break + } + + return cacheMetadata to previousTokenCount + } + + private companion object { + private val logger = LoggerFactory.getLogger(ContextCacheRequestProcessor::class) + } +} diff --git a/core/src/commonTest/kotlin/com/google/adk/kt/processors/ContextCacheRequestProcessorTest.kt b/core/src/commonTest/kotlin/com/google/adk/kt/processors/ContextCacheRequestProcessorTest.kt new file mode 100644 index 00000000..9d5cd273 --- /dev/null +++ b/core/src/commonTest/kotlin/com/google/adk/kt/processors/ContextCacheRequestProcessorTest.kt @@ -0,0 +1,165 @@ +/* + * 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. + */ + +@file:OptIn(com.google.adk.kt.annotations.ExperimentalContextCachingFeature::class) + +package com.google.adk.kt.processors + +import com.google.adk.kt.agents.ContextCacheConfig +import com.google.adk.kt.agents.InvocationContext +import com.google.adk.kt.events.Event +import com.google.adk.kt.models.CacheMetadata +import com.google.adk.kt.models.LlmRequest +import com.google.adk.kt.testing.DummyAgent +import com.google.adk.kt.testing.testSession +import com.google.adk.kt.types.UsageMetadata +import kotlin.test.assertEquals +import kotlin.test.assertNull +import kotlinx.coroutines.runBlocking +import org.junit.Test + +class ContextCacheRequestProcessorTest { + + private val agent = DummyAgent(name = "agent") + + private fun contextWith( + config: ContextCacheConfig?, + events: List = emptyList(), + invocationId: String = "inv-current", + ): InvocationContext { + val session = testSession() + session.events.addAll(events) + return InvocationContext( + session = session, + runConfig = null, + agent = agent, + invocationId = invocationId, + contextCacheConfig = config, + ) + } + + private fun activeCacheMetadata(invocationsUsed: Int) = + CacheMetadata( + fingerprint = "fp", + contentsCount = 3, + cacheName = "cache/1", + expireTime = 1_000L, + invocationsUsed = invocationsUsed, + ) + + @Test + fun process_noCacheConfig_returnsRequestUnchanged() = runBlocking { + val context = contextWith(config = null) + val request = LlmRequest() + + val result = ContextCacheRequestProcessor().process(context, request) + + assertEquals(request, result) + assertNull(result.cacheConfig) + } + + @Test + fun process_cacheConfigButNoEvents_setsConfigOnly() = runBlocking { + val config = ContextCacheConfig() + val context = contextWith(config = config) + + val result = ContextCacheRequestProcessor().process(context, LlmRequest()) + + assertEquals(config, result.cacheConfig) + assertNull(result.cacheMetadata) + assertNull(result.cacheableContentsTokenCount) + } + + @Test + fun process_sameInvocationActiveCache_doesNotIncrementInvocationsUsed() = runBlocking { + val event = + Event( + author = "agent", + invocationId = "inv-current", + cacheMetadata = activeCacheMetadata(invocationsUsed = 2), + usageMetadata = UsageMetadata(promptTokenCount = 5000), + ) + val context = + contextWith( + config = ContextCacheConfig(), + events = listOf(event), + invocationId = "inv-current", + ) + + val result = ContextCacheRequestProcessor().process(context, LlmRequest()) + + assertEquals(2, result.cacheMetadata?.invocationsUsed) + assertEquals(5000, result.cacheableContentsTokenCount) + } + + @Test + fun process_priorInvocationActiveCache_incrementsInvocationsUsed() = runBlocking { + val event = + Event( + author = "agent", + invocationId = "inv-previous", + cacheMetadata = activeCacheMetadata(invocationsUsed = 2), + ) + val context = + contextWith( + config = ContextCacheConfig(), + events = listOf(event), + invocationId = "inv-current", + ) + + val result = ContextCacheRequestProcessor().process(context, LlmRequest()) + + assertEquals(3, result.cacheMetadata?.invocationsUsed) + } + + @Test + fun process_priorInvocationFingerprintOnly_doesNotIncrement() = runBlocking { + val event = + Event( + author = "agent", + invocationId = "inv-previous", + cacheMetadata = CacheMetadata(fingerprint = "fp", contentsCount = 3), + ) + val context = + contextWith( + config = ContextCacheConfig(), + events = listOf(event), + invocationId = "inv-current", + ) + + val result = ContextCacheRequestProcessor().process(context, LlmRequest()) + + assertNull(result.cacheMetadata?.invocationsUsed) + assertEquals(3, result.cacheMetadata?.contentsCount) + } + + @Test + fun process_ignoresEventsFromOtherAgents() = runBlocking { + val otherAgentEvent = + Event( + author = "other", + invocationId = "inv-previous", + cacheMetadata = CacheMetadata(fingerprint = "fp", contentsCount = 1), + usageMetadata = UsageMetadata(promptTokenCount = 9999), + ) + val context = contextWith(config = ContextCacheConfig(), events = listOf(otherAgentEvent)) + + val result = ContextCacheRequestProcessor().process(context, LlmRequest()) + + assertNull(result.cacheMetadata) + assertNull(result.cacheableContentsTokenCount) + } +}