Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -193,6 +194,7 @@ class LlmAgent(
// request context.
CompactionRequestProcessor(),
ContentsProcessor(),
ContextCacheRequestProcessor(),
AgentTransferProcessor(),
OutputSchemaProcessor(),
)
Expand Down
Original file line number Diff line number Diff line change
@@ -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<Event>,
agentName: String,
currentInvocationId: String,
): Pair<CacheMetadata?, Int?> {
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)
}
}
Original file line number Diff line number Diff line change
@@ -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<Event> = 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)
}
}
Loading