diff --git a/core/src/commonJvmAndroidMain/kotlin/com/google/adk/kt/models/Gemini.kt b/core/src/commonJvmAndroidMain/kotlin/com/google/adk/kt/models/Gemini.kt index 3c621210..95554edb 100644 --- a/core/src/commonJvmAndroidMain/kotlin/com/google/adk/kt/models/Gemini.kt +++ b/core/src/commonJvmAndroidMain/kotlin/com/google/adk/kt/models/Gemini.kt @@ -128,11 +128,28 @@ class Gemini( override fun generateContent(request: LlmRequest, stream: Boolean): Flow = flow { val preparedRequest = request.prepareGenerateContentRequest(!client.vertexAI()) - val config = preparedRequest.config.toGenaiSdk() - val contents = preparedRequest.contents.map { it.toGenaiSdk() } + + // Handle context caching when configured. The manager may rewrite the request to reference an + // existing cache (dropping the cached prefix) and returns the metadata to attach to responses. + val cacheManager = + preparedRequest.cacheConfig?.let { + GeminiContextCacheManager(name, GeminiContextCacheManager.RealCacheClient(client.caches)) + } + val cacheResult = cacheManager?.handleContextCaching(preparedRequest) + val finalRequest = cacheResult?.request ?: preparedRequest + val cacheMetadata = cacheResult?.cacheMetadata + + fun attachCacheMetadata(response: LlmResponse): LlmResponse { + // A non-null cacheMetadata implies a non-null cacheManager (both come from cacheResult). + val metadata = cacheMetadata ?: return response + return cacheManager.populateCacheMetadataInResponse(response, metadata) + } + + val config = finalRequest.config.toGenaiSdk() + val contents = finalRequest.contents.map { it.toGenaiSdk() } logger.debug { - "LLM Request:\n${Json.toJsonString(buildLoggingRequestMap(preparedRequest, config))}" + "LLM Request:\n${Json.toJsonString(buildLoggingRequestMap(finalRequest, config))}" } if (stream) { @@ -148,7 +165,7 @@ class Gemini( } // After stream loop ends, emit final aggregated response - aggregator.aggregate()?.let { emit(it) } + aggregator.aggregate()?.let { emit(attachCacheMetadata(it)) } } else { val response = models.generateContent(name, contents, config) logger.debug { @@ -156,7 +173,7 @@ class Gemini( "finishReason=${response.finishReason()}" } val llmResponse = LlmResponse.from(response.fromGenaiSdk()) - emit(llmResponse) + emit(attachCacheMetadata(llmResponse)) } } diff --git a/core/src/commonJvmAndroidMain/kotlin/com/google/adk/kt/models/GeminiContextCacheManager.kt b/core/src/commonJvmAndroidMain/kotlin/com/google/adk/kt/models/GeminiContextCacheManager.kt new file mode 100644 index 00000000..cb7e4c86 --- /dev/null +++ b/core/src/commonJvmAndroidMain/kotlin/com/google/adk/kt/models/GeminiContextCacheManager.kt @@ -0,0 +1,300 @@ +/* + * 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.models + +import com.google.adk.kt.logging.LoggerFactory +import com.google.adk.kt.types.Content +import com.google.adk.kt.types.Role +import com.google.adk.kt.types.toGenaiSdk +import com.google.genai.types.CachedContent +import com.google.genai.types.CreateCachedContentConfig +import java.security.MessageDigest +import kotlin.jvm.optionals.getOrNull +import kotlin.time.Clock + +/** + * Manages the context-cache lifecycle for Gemini models. + * + * This handles cache creation, validation, cleanup, and metadata population for Gemini context + * caching. It uses content hashing (a fingerprint) to determine cache compatibility and applies an + * existing cache to the request by referencing it via [GenerateContentConfig.cachedContent] and + * dropping the cached prefix from the request. + * + * Cache operations are delegated to a [CacheClient] so they can be faked in tests. The Gemini model + * wires in a [RealCacheClient] backed by the GenAI SDK's `client.caches`. + * + * @param modelName The model resource name used when creating caches. + * @param cacheClient The client used to create and delete Gemini caches. + */ +internal class GeminiContextCacheManager( + private val modelName: String, + private val cacheClient: CacheClient, +) { + + /** + * The outcome of [handleContextCaching]: the (possibly rewritten) [request] to send to the model + * and the [cacheMetadata] to attach to the response, or `null` metadata if caching produced no + * usable state. + */ + data class CacheResult(val request: LlmRequest, val cacheMetadata: CacheMetadata?) + + /** + * Validates an existing cache or creates a new one if needed, then applies the cache to the + * request by setting `cachedContent` and dropping the cached prefix. + * + * @param request The request that may carry cache config and metadata. + * @return The (possibly rewritten) request plus the cache metadata to include in the response. + */ + fun handleContextCaching(request: LlmRequest): CacheResult { + val existing = request.cacheMetadata + if (existing != null) { + if (isCacheValid(request)) { + logger.debug { "Cache is valid, reusing cache: ${existing.cacheName}" } + val applied = applyCacheToRequest(request, existing.cacheName!!, existing.contentsCount) + return CacheResult(applied, existing) + } + + // Invalid cache: clean up any active cache, then decide whether to recreate. + if (existing.cacheName != null) { + logger.debug { "Cache is invalid, cleaning up: ${existing.cacheName}" } + cleanupCache(existing.cacheName) + } + + val cacheContentsCount = existing.contentsCount + val currentFingerprint = generateFingerprint(request, cacheContentsCount) + + if (currentFingerprint == existing.fingerprint) { + // Same content but the cache expired/aged out: recreate it. + val created = createNewCacheWithContents(request, cacheContentsCount) + if (created != null) { + val applied = applyCacheToRequest(request, created.cacheName!!, cacheContentsCount) + return CacheResult(applied, created) + } + // Creation failed (e.g. below the token minimum): preserve the prefix fingerprint so the + // fingerprint chain stays stable for subsequent calls. + return CacheResult( + request, + CacheMetadata(fingerprint = currentFingerprint, contentsCount = cacheContentsCount), + ) + } + + // Fingerprints differ: recompute over the current cacheable prefix and return + // fingerprint-only metadata. + val newCount = findCountOfContentsToCache(request.contents) + val fingerprint = generateFingerprint(request, newCount) + return CacheResult( + request, + CacheMetadata(fingerprint = fingerprint, contentsCount = newCount), + ) + } + + // No existing metadata: return fingerprint-only metadata. A cache is never created without a + // prior fingerprint to match against. + val cacheContentsCount = findCountOfContentsToCache(request.contents) + val fingerprint = generateFingerprint(request, cacheContentsCount) + return CacheResult( + request, + CacheMetadata(fingerprint = fingerprint, contentsCount = cacheContentsCount), + ) + } + + /** Returns [response] with [cacheMetadata] attached. */ + fun populateCacheMetadataInResponse( + response: LlmResponse, + cacheMetadata: CacheMetadata, + ): LlmResponse = response.copy(cacheMetadata = cacheMetadata) + + /** + * Finds the number of leading contents to cache: everything before the last contiguous batch of + * user contents. This always leaves at least the latest user turn to send to the API. + */ + private fun findCountOfContentsToCache(contents: List): Int { + if (contents.isEmpty()) return 0 + var lastUserBatchStart = contents.size + for (i in contents.indices.reversed()) { + if (contents[i].role == Role.USER) { + lastUserBatchStart = i + } else { + break + } + } + return lastUserBatchStart + } + + /** Checks whether the cache referenced by the request's metadata is still usable. */ + private fun isCacheValid(request: LlmRequest): Boolean { + val metadata = request.cacheMetadata ?: return false + // Fingerprint-only metadata is not an active cache. + if (metadata.cacheName == null) return false + if (Clock.System.now().toEpochMilliseconds() >= metadata.expireTime!!) { + logger.info { "Cache expired: ${metadata.cacheName}" } + return false + } + val maxInvocations = request.cacheConfig?.maxInvocations ?: return false + if (metadata.invocationsUsed!! > maxInvocations) { + logger.info { "Cache exceeded max invocations: ${metadata.cacheName}" } + return false + } + val currentFingerprint = generateFingerprint(request, metadata.contentsCount) + if (currentFingerprint != metadata.fingerprint) { + logger.debug { "Cache content fingerprint mismatch" } + return false + } + return true + } + + /** + * Generates a 16-character fingerprint over the cacheable state: system instruction, tools, and + * the first [cacheContentsCount] contents. Uses the GenAI SDK JSON form for stable hashing of + * binary parts. + */ + private fun generateFingerprint(request: LlmRequest, cacheContentsCount: Int): String { + val builder = StringBuilder() + request.config.systemInstruction?.let { + builder.append("system_instruction:").append(it.toGenaiSdk().toJson()) + } + request.config.tools?.let { tools -> + builder.append("tools:") + for (tool in tools) { + builder.append(tool.toGenaiSdk().toJson()) + } + } + if (cacheContentsCount > 0 && request.contents.isNotEmpty()) { + builder.append("contents:") + val count = minOf(cacheContentsCount, request.contents.size) + for (i in 0 until count) { + builder.append(request.contents[i].toGenaiSdk().toJson()) + } + } + return sha256Hex(builder.toString()).take(FINGERPRINT_LENGTH) + } + + /** Creates a new cache when the previous request was large enough, otherwise returns `null`. */ + private fun createNewCacheWithContents( + request: LlmRequest, + cacheContentsCount: Int, + ): CacheMetadata? { + val tokenCount = request.cacheableContentsTokenCount + if (tokenCount == null) { + logger.info { "No previous token count available; skipping cache creation." } + return null + } + val minTokens = request.cacheConfig?.minTokens ?: return null + if (tokenCount < minTokens) { + logger.info { "Previous request too small for caching ($tokenCount < $minTokens tokens)." } + return null + } + if (tokenCount < GEMINI_MIN_CACHE_TOKENS) { + logger.info { + "Request below Gemini minimum cache size ($tokenCount < $GEMINI_MIN_CACHE_TOKENS tokens)." + } + return null + } + return try { + createGeminiCache(request, cacheContentsCount) + } catch (e: Exception) { + logger.warn { "Failed to create cache: ${e.message}" } + null + } + } + + /** Creates the cache via the GenAI SDK and returns its metadata. */ + private fun createGeminiCache(request: LlmRequest, cacheContentsCount: Int): CacheMetadata { + val cacheConfig = requireNotNull(request.cacheConfig) { "cacheConfig must be set." } + val cacheContents = request.contents.take(cacheContentsCount).map { it.toGenaiSdk() } + val displayName = "adk-cache-${Clock.System.now().epochSeconds}-${cacheContentsCount}contents" + + val configBuilder = + CreateCachedContentConfig.builder() + .contents(cacheContents) + .ttl(java.time.Duration.ofMillis(cacheConfig.ttl.inWholeMilliseconds)) + .displayName(displayName) + request.config.systemInstruction?.let { configBuilder.systemInstruction(it.toGenaiSdk()) } + request.config.tools?.let { tools -> configBuilder.tools(tools.map { it.toGenaiSdk() }) } + + val cachedContent: CachedContent = cacheClient.create(modelName, configBuilder.build()) + val createdAt = Clock.System.now().toEpochMilliseconds() + val cacheName = + cachedContent.name().getOrNull() + ?: throw IllegalStateException("Created cache has no resource name.") + logger.info { "Cache created successfully: $cacheName" } + + return CacheMetadata( + fingerprint = generateFingerprint(request, cacheContentsCount), + contentsCount = cacheContentsCount, + cacheName = cacheName, + expireTime = createdAt + cacheConfig.ttl.inWholeMilliseconds, + invocationsUsed = 1, + createdAt = createdAt, + ) + } + + /** Deletes a cache, logging and swallowing any failure. */ + private fun cleanupCache(cacheName: String) { + try { + cacheClient.delete(cacheName) + logger.info { "Cache cleaned up: $cacheName" } + } catch (e: Exception) { + logger.warn { "Failed to cleanup cache $cacheName: ${e.message}" } + } + } + + /** + * Rewrites the request to use the cache: references it via `cachedContent`, drops the cached + * system instruction and tools, and removes the cached content prefix. + */ + private fun applyCacheToRequest( + request: LlmRequest, + cacheName: String, + cacheContentsCount: Int, + ): LlmRequest = + request.copy( + config = + request.config.copy(systemInstruction = null, tools = null, cachedContent = cacheName), + contents = request.contents.drop(cacheContentsCount), + ) + + private fun sha256Hex(input: String): String { + val bytes = MessageDigest.getInstance("SHA-256").digest(input.encodeToByteArray()) + return bytes.joinToString("") { (it.toInt() and 0xff).toString(16).padStart(2, '0') } + } + + /** Abstraction over the GenAI SDK cache operations to allow faking in tests. */ + interface CacheClient { + fun create(model: String, config: CreateCachedContentConfig): CachedContent + + fun delete(name: String) + } + + /** [CacheClient] backed by the GenAI SDK's `client.caches`. */ + class RealCacheClient(private val caches: com.google.genai.Caches) : CacheClient { + override fun create(model: String, config: CreateCachedContentConfig): CachedContent = + caches.create(model, config) + + override fun delete(name: String) { + val unused = caches.delete(name, null) + } + } + + private companion object { + private val logger = LoggerFactory.getLogger(GeminiContextCacheManager::class) + + // Gemini requires a minimum number of tokens before content can be cached. + private const val GEMINI_MIN_CACHE_TOKENS = 4096 + + private const val FINGERPRINT_LENGTH = 16 + } +} diff --git a/core/src/commonJvmAndroidMain/kotlin/com/google/adk/kt/types/GenaiConverters.kt b/core/src/commonJvmAndroidMain/kotlin/com/google/adk/kt/types/GenaiConverters.kt index 571135ae..9dac7cf1 100644 --- a/core/src/commonJvmAndroidMain/kotlin/com/google/adk/kt/types/GenaiConverters.kt +++ b/core/src/commonJvmAndroidMain/kotlin/com/google/adk/kt/types/GenaiConverters.kt @@ -224,6 +224,7 @@ internal fun com.google.genai.types.GenerateContentConfig.fromGenaiSdk(): Genera responseMimeType = responseMimeType().getOrNull(), responseSchema = responseSchema().getOrNull()?.toKtSchema(), thinkingConfig = thinkingConfig().getOrNull()?.fromGenaiSdk(), + cachedContent = cachedContent().getOrNull(), ) /** @@ -245,6 +246,7 @@ internal fun GenerateContentConfig.toGenaiSdk(): com.google.genai.types.Generate this@toGenaiSdk.responseMimeType?.let { responseMimeType(it) } this@toGenaiSdk.responseSchema?.let { responseSchema(it.toGenAiSchema()) } this@toGenaiSdk.thinkingConfig?.let { thinkingConfig(it.toGenaiSdk()) } + this@toGenaiSdk.cachedContent?.let { cachedContent(it) } } .build() diff --git a/core/src/commonJvmAndroidTest/kotlin/com/google/adk/kt/types/GenaiConvertersTest.kt b/core/src/commonJvmAndroidTest/kotlin/com/google/adk/kt/types/GenaiConvertersTest.kt index 360ad809..02ae0985 100644 --- a/core/src/commonJvmAndroidTest/kotlin/com/google/adk/kt/types/GenaiConvertersTest.kt +++ b/core/src/commonJvmAndroidTest/kotlin/com/google/adk/kt/types/GenaiConvertersTest.kt @@ -256,6 +256,17 @@ class GenaiConvertersTest { assertEquals(40, adkConfig.topK) } + @Test + fun generateContentConfig_cachedContentRoundTrip_preservesValue() { + val cacheName = "projects/p/locations/l/cachedContents/456" + val adkConfig = GenerateContentConfig(cachedContent = cacheName) + + val genaiConfig = adkConfig.toGenaiSdk() + + assertEquals(cacheName, genaiConfig.cachedContent().get()) + assertEquals(cacheName, genaiConfig.fromGenaiSdk().cachedContent) + } + @Test fun thinkingConfig_convertsCorrectly() { val adkThinkingConfig = diff --git a/core/src/commonMain/kotlin/com/google/adk/kt/agents/ContextCacheConfig.kt b/core/src/commonMain/kotlin/com/google/adk/kt/agents/ContextCacheConfig.kt new file mode 100644 index 00000000..87bf01e1 --- /dev/null +++ b/core/src/commonMain/kotlin/com/google/adk/kt/agents/ContextCacheConfig.kt @@ -0,0 +1,56 @@ +/* + * 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.agents + +import kotlin.time.Duration +import kotlin.time.Duration.Companion.seconds + +/** + * Configuration for context caching across all agents in an app. + * + * This configuration enables and controls context caching behavior for all LLM agents in an app. + * When this config is present on an app, context caching is enabled for all agents. When absent + * (`null`), context caching is disabled. + * + * Context caching can significantly reduce costs and improve response times by reusing previously + * processed context across multiple requests. + * + * @property maxInvocations Maximum number of invocations to reuse the same cache before refreshing + * it. Must be in the range `1..100`. Defaults to 10. + * @property ttl Time-to-live for the cache. Must be strictly positive. Defaults to 30 minutes. + * @property minTokens Minimum estimated request tokens required to enable caching. This compares + * against the estimated total tokens of the request (system instruction + tools + contents). + * Context cache storage may have a cost, so set this higher to avoid caching small requests where + * the overhead may exceed the benefits. Must be non-negative. Defaults to 0. + */ +data class ContextCacheConfig( + val maxInvocations: Int = 10, + val ttl: Duration = 1800.seconds, + val minTokens: Int = 0, +) { + init { + require(maxInvocations in 1..100) { + "maxInvocations must be in 1..100, but was $maxInvocations." + } + require(ttl.isPositive()) { "ttl must be positive, but was $ttl." } + require(minTokens >= 0) { "minTokens must be >= 0, but was $minTokens." } + } + + /** Returns the TTL in the string format used for cache creation, e.g. `"1800s"`. */ + val ttlString: String + get() = "${ttl.inWholeSeconds}s" +} diff --git a/core/src/commonMain/kotlin/com/google/adk/kt/agents/InvocationContext.kt b/core/src/commonMain/kotlin/com/google/adk/kt/agents/InvocationContext.kt index 7f83ed4c..b368c016 100644 --- a/core/src/commonMain/kotlin/com/google/adk/kt/agents/InvocationContext.kt +++ b/core/src/commonMain/kotlin/com/google/adk/kt/agents/InvocationContext.kt @@ -121,6 +121,12 @@ data class InvocationContext( */ val resumabilityConfig: ResumabilityConfig? = null, + /** + * Optional context cache configuration for this invocation, propagated from the [App]. When + * `null`, context caching is disabled for the invocation. + */ + val contextCacheConfig: ContextCacheConfig? = null, + // State /** The user content that started this invocation. Readonly. */ val userContent: Content? = null, 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 af72fa3d..4a4651da 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 @@ -31,6 +31,7 @@ import com.google.adk.kt.models.Model import com.google.adk.kt.processors.AgentTransferProcessor import com.google.adk.kt.processors.BasicRequestProcessor 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 @@ -189,6 +190,7 @@ class LlmAgent( RequestConfirmationProcessor(), InstructionsProcessor(), ContentsProcessor(), + ContextCacheRequestProcessor(), AgentTransferProcessor(), OutputSchemaProcessor(), ) diff --git a/core/src/commonMain/kotlin/com/google/adk/kt/agents/LlmAgentTurn.kt b/core/src/commonMain/kotlin/com/google/adk/kt/agents/LlmAgentTurn.kt index b501d32f..55b1bfe3 100644 --- a/core/src/commonMain/kotlin/com/google/adk/kt/agents/LlmAgentTurn.kt +++ b/core/src/commonMain/kotlin/com/google/adk/kt/agents/LlmAgentTurn.kt @@ -345,6 +345,7 @@ internal class LlmAgentTurn( errorMessage = response.errorMessage, partial = response.partial, interrupted = response.interrupted, + cacheMetadata = response.cacheMetadata, ) .populateClientFunctionCallId() diff --git a/core/src/commonMain/kotlin/com/google/adk/kt/apps/App.kt b/core/src/commonMain/kotlin/com/google/adk/kt/apps/App.kt index 3d9e86e8..190e52b2 100644 --- a/core/src/commonMain/kotlin/com/google/adk/kt/apps/App.kt +++ b/core/src/commonMain/kotlin/com/google/adk/kt/apps/App.kt @@ -17,6 +17,7 @@ package com.google.adk.kt.apps import com.google.adk.kt.agents.BaseAgent +import com.google.adk.kt.agents.ContextCacheConfig import com.google.adk.kt.agents.ResumabilityConfig import com.google.adk.kt.plugins.Plugin import com.google.adk.kt.summarizer.EventsCompactionConfig @@ -42,6 +43,8 @@ import com.google.adk.kt.summarizer.EventsCompactionConfig * sessions. When `null`, resumability is disabled. * @property eventsCompactionConfig Optional configuration controlling context-compaction strategies * for sessions of this application. When `null`, no compaction runs. + * @property contextCacheConfig Optional context cache configuration that applies to all LLM agents + * in the app. When `null`, context caching is disabled. */ data class App( val appName: String, @@ -49,6 +52,7 @@ data class App( val plugins: List = emptyList(), val resumabilityConfig: ResumabilityConfig? = null, val eventsCompactionConfig: EventsCompactionConfig? = null, + val contextCacheConfig: ContextCacheConfig? = null, ) { init { require(IDENTIFIER_REGEX.matches(appName)) { diff --git a/core/src/commonMain/kotlin/com/google/adk/kt/events/Event.kt b/core/src/commonMain/kotlin/com/google/adk/kt/events/Event.kt index 20c79579..fa08dd52 100644 --- a/core/src/commonMain/kotlin/com/google/adk/kt/events/Event.kt +++ b/core/src/commonMain/kotlin/com/google/adk/kt/events/Event.kt @@ -17,6 +17,7 @@ package com.google.adk.kt.events import com.google.adk.kt.ids.Uuid +import com.google.adk.kt.models.CacheMetadata import com.google.adk.kt.tools.BaseTool import com.google.adk.kt.types.CitationMetadata import com.google.adk.kt.types.Content @@ -53,6 +54,8 @@ import kotlinx.serialization.Serializable * multiple sub-agents shouldn't see their peer agents' conversation history. * @property groundingMetadata The grounding metadata of the event. * @property modelVersion The model version used to generate the response. + * @property cacheMetadata Context cache metadata associated with this event's LLM response, used to + * carry cache state across turns. `null` when context caching is disabled. * @property timestamp The timestamp of the event. */ @Serializable @@ -75,6 +78,7 @@ data class Event( val groundingMetadata: GroundingMetadata? = null, val modelVersion: String? = null, val citationMetadata: CitationMetadata? = null, + val cacheMetadata: CacheMetadata? = null, val customMetadata: Map? = null, val timestamp: Long = Clock.System.now().toEpochMilliseconds(), ) { diff --git a/core/src/commonMain/kotlin/com/google/adk/kt/models/CacheMetadata.kt b/core/src/commonMain/kotlin/com/google/adk/kt/models/CacheMetadata.kt new file mode 100644 index 00000000..3b156f41 --- /dev/null +++ b/core/src/commonMain/kotlin/com/google/adk/kt/models/CacheMetadata.kt @@ -0,0 +1,95 @@ +/* + * 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.models + +import kotlin.time.Clock +import kotlin.time.Duration.Companion.minutes +import kotlinx.serialization.Serializable + +/** + * Metadata for the context cache associated with LLM responses. + * + * This stores cache identification, usage tracking, and lifecycle information for a particular + * cache instance. It can be in one of two states: + * 1. **Active cache state:** [cacheName] is set and [expireTime], [invocationsUsed], and + * [createdAt] are populated. + * 2. **Fingerprint-only state:** [cacheName] is `null` and only [fingerprint] and [contentsCount] + * are set, used for prefix matching before a cache exists. + * + * Token counts (cached and total) are available in [LlmResponse.usageMetadata] and should be + * accessed from there to avoid duplication. + * + * @property fingerprint Hash of the cacheable contents (system instruction + tools + contents). + * Always present for prefix matching. + * @property contentsCount Number of contents. When an active cache exists this is the count of + * cached contents; otherwise it is the count of the cacheable content prefix used for + * fingerprinting. Must be non-negative. + * @property cacheName Full resource name of the cached content (e.g. + * `projects/123/locations/us-central1/cachedContents/456`). `null` when no active cache exists. + * @property expireTime Epoch milliseconds when the cache expires. `null` when no active cache + * exists. + * @property invocationsUsed Number of invocations this cache has been used for. `null` when no + * active cache exists. Must be non-negative when set. + * @property createdAt Epoch milliseconds when the cache was created. `null` when no active cache + * exists. + */ +@Serializable +data class CacheMetadata( + val fingerprint: String, + val contentsCount: Int, + val cacheName: String? = null, + val expireTime: Long? = null, + val invocationsUsed: Int? = null, + val createdAt: Long? = null, +) { + init { + require(contentsCount >= 0) { "contentsCount must be >= 0, but was $contentsCount." } + invocationsUsed?.let { require(it >= 0) { "invocationsUsed must be >= 0, but was $it." } } + val activeFlags = listOf(cacheName != null, expireTime != null, invocationsUsed != null) + require(activeFlags.all { it } || activeFlags.none { it }) { + "cacheName, expireTime, and invocationsUsed must all be set (active cache) or all be null " + + "(fingerprint-only state)." + } + } + + /** Whether this metadata refers to an active cache, as opposed to a fingerprint-only state. */ + val isActive: Boolean + get() = cacheName != null + + /** + * Whether the cache will expire within [EXPIRY_BUFFER]. Always `false` for fingerprint-only + * metadata (which has no [expireTime]). + */ + val expireSoon: Boolean + get() { + val expiry = expireTime ?: return false + return Clock.System.now().toEpochMilliseconds() > expiry - EXPIRY_BUFFER.inWholeMilliseconds + } + + override fun toString(): String { + val name = + cacheName + ?: return "Fingerprint-only: $contentsCount contents, " + + "fingerprint=${fingerprint.take(8)}..." + val cacheId = name.substringAfterLast("/") + return "Cache $cacheId: used $invocationsUsed invocations, cached $contentsCount contents" + } + + private companion object { + // Buffer applied when checking expiry so a cache nearing expiry is treated as expired. + val EXPIRY_BUFFER = 2.minutes + } +} diff --git a/core/src/commonMain/kotlin/com/google/adk/kt/models/LlmRequest.kt b/core/src/commonMain/kotlin/com/google/adk/kt/models/LlmRequest.kt index c001d96b..7ce2f90c 100644 --- a/core/src/commonMain/kotlin/com/google/adk/kt/models/LlmRequest.kt +++ b/core/src/commonMain/kotlin/com/google/adk/kt/models/LlmRequest.kt @@ -15,6 +15,7 @@ */ package com.google.adk.kt.models +import com.google.adk.kt.agents.ContextCacheConfig import com.google.adk.kt.tools.BaseTool import com.google.adk.kt.types.Blob import com.google.adk.kt.types.Content @@ -32,12 +33,21 @@ import com.google.adk.kt.types.Tool * @property contents The contents of the request. * @property config The configuration for generating content. * @property toolsDict Internal mapping of tools added during request processing. + * @property cacheConfig Context cache configuration for this request. When `null`, context caching + * is disabled for the request. + * @property cacheMetadata Cache metadata carried over from previous requests, used to validate and + * reuse an existing cache. + * @property cacheableContentsTokenCount Prompt token count from the previous request, used to gate + * cache creation on a minimum size. */ data class LlmRequest( val model: Model? = null, val contents: List = emptyList(), val config: GenerateContentConfig = GenerateContentConfig(), internal val toolsDict: List = emptyList(), + val cacheConfig: ContextCacheConfig? = null, + val cacheMetadata: CacheMetadata? = null, + val cacheableContentsTokenCount: Int? = null, ) { /** * Appends tools to the request and merges any new function declarations. diff --git a/core/src/commonMain/kotlin/com/google/adk/kt/models/LlmResponse.kt b/core/src/commonMain/kotlin/com/google/adk/kt/models/LlmResponse.kt index 3c08a50f..135ec707 100644 --- a/core/src/commonMain/kotlin/com/google/adk/kt/models/LlmResponse.kt +++ b/core/src/commonMain/kotlin/com/google/adk/kt/models/LlmResponse.kt @@ -38,6 +38,8 @@ import com.google.adk.kt.types.UsageMetadata * @property modelVersion The model version used to generate the response. * @property citationMetadata The citation metadata of the response. * @property groundingMetadata The grounding metadata of the response. + * @property cacheMetadata Context cache metadata for this response, populated when context caching + * is enabled. `null` when caching is disabled or no cache information is available. */ data class LlmResponse( val content: Content? = null, @@ -49,6 +51,7 @@ data class LlmResponse( val modelVersion: String? = null, val citationMetadata: CitationMetadata? = null, val groundingMetadata: GroundingMetadata? = null, + val cacheMetadata: CacheMetadata? = null, ) { companion object { /** 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..cfcff496 --- /dev/null +++ b/core/src/commonMain/kotlin/com/google/adk/kt/processors/ContextCacheRequestProcessor.kt @@ -0,0 +1,98 @@ +/* + * 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.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) + + 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 + } +} diff --git a/core/src/commonMain/kotlin/com/google/adk/kt/runners/AbstractRunner.kt b/core/src/commonMain/kotlin/com/google/adk/kt/runners/AbstractRunner.kt index beb1f76c..b97077b1 100644 --- a/core/src/commonMain/kotlin/com/google/adk/kt/runners/AbstractRunner.kt +++ b/core/src/commonMain/kotlin/com/google/adk/kt/runners/AbstractRunner.kt @@ -19,6 +19,7 @@ package com.google.adk.kt.runners import com.google.adk.kt.agents.BaseAgent +import com.google.adk.kt.agents.ContextCacheConfig import com.google.adk.kt.agents.InvocationContext import com.google.adk.kt.agents.LlmAgent import com.google.adk.kt.agents.ResumabilityConfig @@ -69,6 +70,13 @@ abstract class AbstractRunner : Runner { final override val pluginManager: PluginManager final override val resumabilityConfig: ResumabilityConfig + /** + * The context cache configuration applied to invocations created by this runner, or `null` when + * context caching is disabled. Only the [App]-based constructor populates this from + * [App.contextCacheConfig]. + */ + val contextCacheConfig: ContextCacheConfig? + /** Creates a runner from explicit fields, not using an [App]. */ constructor( appName: String, @@ -87,6 +95,7 @@ abstract class AbstractRunner : Runner { this.pluginManager = pluginManager this.resumabilityConfig = resumabilityConfig this.app = null + this.contextCacheConfig = null } /** @@ -114,6 +123,7 @@ abstract class AbstractRunner : Runner { this.memoryService = memoryService this.pluginManager = PluginManager(app.plugins) this.resumabilityConfig = app.resumabilityConfig ?: ResumabilityConfig() + this.contextCacheConfig = app.contextCacheConfig this.app = app.copy( eventsCompactionConfig = @@ -534,6 +544,7 @@ abstract class AbstractRunner : Runner { userContent = newMessage, pluginManager = pluginManager, resumabilityConfig = resumabilityConfig, + contextCacheConfig = contextCacheConfig, ) .let { // Run callbacks and append user message to session @@ -589,6 +600,7 @@ abstract class AbstractRunner : Runner { userContent = userMessage, pluginManager = pluginManager, resumabilityConfig = resumabilityConfig, + contextCacheConfig = contextCacheConfig, ) val currentContext = diff --git a/core/src/commonMain/kotlin/com/google/adk/kt/types/GenerateContentConfig.kt b/core/src/commonMain/kotlin/com/google/adk/kt/types/GenerateContentConfig.kt index 40432c3e..6fe8c8f5 100644 --- a/core/src/commonMain/kotlin/com/google/adk/kt/types/GenerateContentConfig.kt +++ b/core/src/commonMain/kotlin/com/google/adk/kt/types/GenerateContentConfig.kt @@ -16,7 +16,13 @@ package com.google.adk.kt.types -/** Configuration for generating content. */ +/** + * Configuration for generating content. + * + * @property cachedContent Resource name of a context cache (e.g. + * `projects/123/locations/us-central1/cachedContents/456`) to use for this request. When set, the + * cached system instruction, tools, and contents are reused instead of being resent. + */ data class GenerateContentConfig( val tools: List? = null, val labels: Map? = null, @@ -30,4 +36,5 @@ data class GenerateContentConfig( val responseMimeType: String? = null, val responseSchema: Schema? = null, val thinkingConfig: ThinkingConfig? = null, + val cachedContent: String? = null, ) diff --git a/core/src/commonTest/kotlin/com/google/adk/kt/agents/ContextCacheConfigTest.kt b/core/src/commonTest/kotlin/com/google/adk/kt/agents/ContextCacheConfigTest.kt new file mode 100644 index 00000000..8acac1c3 --- /dev/null +++ b/core/src/commonTest/kotlin/com/google/adk/kt/agents/ContextCacheConfigTest.kt @@ -0,0 +1,73 @@ +/* + * 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.agents + +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.time.Duration.Companion.seconds + +class ContextCacheConfigTest { + + @Test + fun construct_defaults_usesDocumentedValues() { + val config = ContextCacheConfig() + + assertEquals(10, config.maxInvocations) + assertEquals(1800.seconds, config.ttl) + assertEquals(0, config.minTokens) + } + + @Test + fun ttlString_default_formatsSeconds() { + assertEquals("1800s", ContextCacheConfig().ttlString) + } + + @Test + fun ttlString_customTtl_formatsWholeSeconds() { + assertEquals("60s", ContextCacheConfig(ttl = 60.seconds).ttlString) + } + + @Test + fun construct_customValues_exposesProperties() { + val config = ContextCacheConfig(maxInvocations = 5, ttl = 120.seconds, minTokens = 4096) + + assertEquals(5, config.maxInvocations) + assertEquals(120.seconds, config.ttl) + assertEquals(4096, config.minTokens) + } + + @Test + fun construct_maxInvocationsTooLow_throwsIllegalArgumentException() { + assertFailsWith { ContextCacheConfig(maxInvocations = 0) } + } + + @Test + fun construct_maxInvocationsTooHigh_throwsIllegalArgumentException() { + assertFailsWith { ContextCacheConfig(maxInvocations = 101) } + } + + @Test + fun construct_nonPositiveTtl_throwsIllegalArgumentException() { + assertFailsWith { ContextCacheConfig(ttl = 0.seconds) } + } + + @Test + fun construct_negativeMinTokens_throwsIllegalArgumentException() { + assertFailsWith { ContextCacheConfig(minTokens = -1) } + } +} diff --git a/core/src/commonTest/kotlin/com/google/adk/kt/apps/AppTest.kt b/core/src/commonTest/kotlin/com/google/adk/kt/apps/AppTest.kt index 23b6911b..1bf77e41 100644 --- a/core/src/commonTest/kotlin/com/google/adk/kt/apps/AppTest.kt +++ b/core/src/commonTest/kotlin/com/google/adk/kt/apps/AppTest.kt @@ -18,6 +18,7 @@ package com.google.adk.kt.apps +import com.google.adk.kt.agents.ContextCacheConfig import com.google.adk.kt.agents.ResumabilityConfig import com.google.adk.kt.annotations.ExperimentalResumabilityFeature import com.google.adk.kt.plugins.Plugin @@ -57,6 +58,22 @@ class AppTest { assertSame(config, app.eventsCompactionConfig) } + @Test + fun construct_noContextCacheConfig_defaultsToNull() { + val app = App(appName = "my_app", rootAgent = DummyAgent()) + + assertNull(app.contextCacheConfig) + } + + @Test + fun construct_withContextCacheConfig_exposesIt() { + val config = ContextCacheConfig(maxInvocations = 5) + + val app = App(appName = "my_app", rootAgent = DummyAgent(), contextCacheConfig = config) + + assertSame(config, app.contextCacheConfig) + } + @Test fun construct_emptyName_throwsIllegalArgumentException() { assertFailsWith { App(appName = "", rootAgent = DummyAgent()) } diff --git a/core/src/commonTest/kotlin/com/google/adk/kt/models/CacheMetadataTest.kt b/core/src/commonTest/kotlin/com/google/adk/kt/models/CacheMetadataTest.kt new file mode 100644 index 00000000..303cc291 --- /dev/null +++ b/core/src/commonTest/kotlin/com/google/adk/kt/models/CacheMetadataTest.kt @@ -0,0 +1,123 @@ +/* + * 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.models + +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertFalse +import kotlin.test.assertTrue +import kotlin.time.Clock + +class CacheMetadataTest { + + @Test + fun construct_fingerprintOnly_isNotActive() { + val metadata = CacheMetadata(fingerprint = "abc123", contentsCount = 3) + + assertFalse(metadata.isActive) + assertFalse(metadata.expireSoon) + } + + @Test + fun construct_activeCache_exposesProperties() { + val metadata = + CacheMetadata( + fingerprint = "abc123", + contentsCount = 3, + cacheName = "projects/p/locations/l/cachedContents/456", + expireTime = Clock.System.now().toEpochMilliseconds() + 1_800_000, + invocationsUsed = 1, + createdAt = Clock.System.now().toEpochMilliseconds(), + ) + + assertTrue(metadata.isActive) + assertEquals(1, metadata.invocationsUsed) + } + + @Test + fun construct_partialActiveState_throwsIllegalArgumentException() { + // cacheName set but expireTime/invocationsUsed null violates the active-state invariant. + assertFailsWith { + CacheMetadata(fingerprint = "abc", contentsCount = 1, cacheName = "cache/1") + } + } + + @Test + fun construct_negativeContentsCount_throwsIllegalArgumentException() { + assertFailsWith { + CacheMetadata(fingerprint = "abc", contentsCount = -1) + } + } + + @Test + fun construct_negativeInvocationsUsed_throwsIllegalArgumentException() { + assertFailsWith { + CacheMetadata( + fingerprint = "abc", + contentsCount = 1, + cacheName = "cache/1", + expireTime = 1L, + invocationsUsed = -1, + ) + } + } + + @Test + fun expireSoon_pastExpiry_isTrue() { + val metadata = + CacheMetadata( + fingerprint = "abc", + contentsCount = 1, + cacheName = "cache/1", + expireTime = Clock.System.now().toEpochMilliseconds() - 1_000, + invocationsUsed = 1, + ) + + assertTrue(metadata.expireSoon) + } + + @Test + fun expireSoon_farFutureExpiry_isFalse() { + val metadata = + CacheMetadata( + fingerprint = "abc", + contentsCount = 1, + cacheName = "cache/1", + expireTime = Clock.System.now().toEpochMilliseconds() + 3_600_000, + invocationsUsed = 1, + ) + + assertFalse(metadata.expireSoon) + } + + @Test + fun copy_incrementInvocationsUsed_producesUpdatedCopy() { + val metadata = + CacheMetadata( + fingerprint = "abc", + contentsCount = 1, + cacheName = "cache/1", + expireTime = Clock.System.now().toEpochMilliseconds() + 1_000, + invocationsUsed = 1, + ) + + val updated = metadata.copy(invocationsUsed = metadata.invocationsUsed!! + 1) + + assertEquals(2, updated.invocationsUsed) + assertEquals(1, metadata.invocationsUsed) + } +} diff --git a/core/src/commonTest/kotlin/com/google/adk/kt/models/LlmRequestTest.kt b/core/src/commonTest/kotlin/com/google/adk/kt/models/LlmRequestTest.kt index 48aa4ad6..8983348d 100644 --- a/core/src/commonTest/kotlin/com/google/adk/kt/models/LlmRequestTest.kt +++ b/core/src/commonTest/kotlin/com/google/adk/kt/models/LlmRequestTest.kt @@ -16,6 +16,7 @@ package com.google.adk.kt.models +import com.google.adk.kt.agents.ContextCacheConfig import com.google.adk.kt.testing.userMessage import com.google.adk.kt.tools.BaseTool import com.google.adk.kt.tools.ToolContext @@ -28,6 +29,7 @@ import com.google.adk.kt.types.Role import kotlin.test.Test import kotlin.test.assertEquals import kotlin.test.assertNotNull +import kotlin.test.assertNull import kotlin.test.assertTrue class LlmRequestTest { @@ -198,4 +200,26 @@ class LlmRequestTest { val sysText = systemInstruction.parts.joinToString("") { it.text ?: "" } assertEquals("Initial\n\nAppended", sysText) } + + @Test + fun cacheFields_defaultToNull() { + val request = LlmRequest() + + assertNull(request.cacheConfig) + assertNull(request.cacheMetadata) + assertNull(request.cacheableContentsTokenCount) + } + + @Test + fun cacheFields_canBeSet() { + val config = ContextCacheConfig() + val metadata = CacheMetadata(fingerprint = "abc", contentsCount = 2) + + val request = + LlmRequest(cacheConfig = config, cacheMetadata = metadata, cacheableContentsTokenCount = 5000) + + assertEquals(config, request.cacheConfig) + assertEquals(metadata, request.cacheMetadata) + assertEquals(5000, request.cacheableContentsTokenCount) + } } diff --git a/core/src/commonTest/kotlin/com/google/adk/kt/models/LlmResponseTest.kt b/core/src/commonTest/kotlin/com/google/adk/kt/models/LlmResponseTest.kt index 8806e9d7..408ed892 100644 --- a/core/src/commonTest/kotlin/com/google/adk/kt/models/LlmResponseTest.kt +++ b/core/src/commonTest/kotlin/com/google/adk/kt/models/LlmResponseTest.kt @@ -101,4 +101,18 @@ class LlmResponseTest { val llmResponse = LlmResponse.from(response) assertEquals("gemini-2.0-flash", llmResponse.modelVersion) } + + @Test + fun cacheMetadata_defaultsToNull() { + assertEquals(null, LlmResponse().cacheMetadata) + } + + @Test + fun cacheMetadata_canBeSet() { + val metadata = CacheMetadata(fingerprint = "abc", contentsCount = 2) + + val llmResponse = LlmResponse(cacheMetadata = metadata) + + assertEquals(metadata, llmResponse.cacheMetadata) + } } 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..1f9d31af --- /dev/null +++ b/core/src/commonTest/kotlin/com/google/adk/kt/processors/ContextCacheRequestProcessorTest.kt @@ -0,0 +1,163 @@ +/* + * 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.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) + } +} diff --git a/core/src/commonTest/kotlin/com/google/adk/kt/runners/AbstractRunnerTest.kt b/core/src/commonTest/kotlin/com/google/adk/kt/runners/AbstractRunnerTest.kt index 24399396..e3cdc269 100644 --- a/core/src/commonTest/kotlin/com/google/adk/kt/runners/AbstractRunnerTest.kt +++ b/core/src/commonTest/kotlin/com/google/adk/kt/runners/AbstractRunnerTest.kt @@ -19,6 +19,7 @@ package com.google.adk.kt.runners import com.google.adk.kt.agents.BaseAgent +import com.google.adk.kt.agents.ContextCacheConfig import com.google.adk.kt.agents.InvocationContext import com.google.adk.kt.agents.LlmAgent import com.google.adk.kt.agents.ResumabilityConfig @@ -58,8 +59,11 @@ class AbstractRunnerTest { class TestRunner(agent: BaseAgent, resumable: Boolean = true) : InMemoryRunner( - agent = agent, - resumabilityConfig = ResumabilityConfig(isResumable = resumable), + App( + appName = "InMemoryRunner", + rootAgent = agent, + resumabilityConfig = ResumabilityConfig(isResumable = resumable), + ) ) { suspend fun callFindAgentToRun(context: InvocationContext, rootAgent: BaseAgent): BaseAgent { return findAgentToRun(context, rootAgent) @@ -812,6 +816,55 @@ class AbstractRunnerTest { } } + // ----- Context cache config propagation ----- + + @Test + fun runAsync_appWithContextCacheConfig_propagatesToInvocationContext() = runTest { + val cacheConfig = ContextCacheConfig(maxInvocations = 5) + var capturedConfig: ContextCacheConfig? = null + val agent = + DummyAgent(name = "agent") { context -> + capturedConfig = context.contextCacheConfig + emit( + Event( + author = Role.MODEL, + invocationId = context.invocationId, + content = modelMessage("resp"), + ) + ) + } + val runner = + InMemoryRunner( + app = App(appName = "test_app", rootAgent = agent, contextCacheConfig = cacheConfig) + ) + + runner.runAsync(userId = "user", sessionId = "session", newMessage = userMessage("hi")).toList() + + assertEquals(cacheConfig, capturedConfig) + } + + @Test + fun runAsync_appWithoutContextCacheConfig_invocationContextConfigIsNull() = runTest { + // Sentinel non-null so the assertion meaningfully verifies the runner left it unset. + var capturedConfig: ContextCacheConfig? = ContextCacheConfig() + val agent = + DummyAgent(name = "agent") { context -> + capturedConfig = context.contextCacheConfig + emit( + Event( + author = Role.MODEL, + invocationId = context.invocationId, + content = modelMessage("resp"), + ) + ) + } + val runner = InMemoryRunner(app = App(appName = "test_app", rootAgent = agent)) + + runner.runAsync(userId = "user", sessionId = "session", newMessage = userMessage("hi")).toList() + + assertNull(capturedConfig) + } + @Test fun runAsync_defaultSummarizerUsesRootLlmAgentModel() = runTest { // No summarizer supplied: the runner must build a default LlmEventSummarizer from the root diff --git a/core/src/commonTest/kotlin/com/google/adk/kt/serialization/EventSerializationTest.kt b/core/src/commonTest/kotlin/com/google/adk/kt/serialization/EventSerializationTest.kt index 01998ba1..a9469794 100644 --- a/core/src/commonTest/kotlin/com/google/adk/kt/serialization/EventSerializationTest.kt +++ b/core/src/commonTest/kotlin/com/google/adk/kt/serialization/EventSerializationTest.kt @@ -21,6 +21,7 @@ import com.google.adk.kt.annotations.FrameworkInternalApi import com.google.adk.kt.events.Event import com.google.adk.kt.events.EventActions import com.google.adk.kt.events.EventCompaction +import com.google.adk.kt.models.CacheMetadata import com.google.adk.kt.sessions.State import com.google.adk.kt.types.Blob import com.google.adk.kt.types.Content @@ -97,6 +98,26 @@ class EventSerializationTest { assertEquals(event, roundTrip(event)) } + @Test + fun cacheMetadata_activeCache_roundTripsLosslessly() { + val event = + Event( + author = "agent", + cacheMetadata = + CacheMetadata( + fingerprint = "abc123", + contentsCount = 4, + cacheName = "projects/p/locations/l/cachedContents/456", + expireTime = 1_800_000L, + invocationsUsed = 2, + createdAt = 1_000L, + ), + timestamp = 1L, + ) + + assertEquals(event.cacheMetadata, roundTrip(event).cacheMetadata) + } + @Test fun stateDelta_removedSentinel_roundTripsToRemoved() { val event = diff --git a/core/src/jvmTest/kotlin/com/google/adk/kt/models/GeminiContextCacheManagerTest.kt b/core/src/jvmTest/kotlin/com/google/adk/kt/models/GeminiContextCacheManagerTest.kt new file mode 100644 index 00000000..bfcbbb1e --- /dev/null +++ b/core/src/jvmTest/kotlin/com/google/adk/kt/models/GeminiContextCacheManagerTest.kt @@ -0,0 +1,170 @@ +/* + * 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.models + +import com.google.adk.kt.agents.ContextCacheConfig +import com.google.adk.kt.types.Content +import com.google.adk.kt.types.GenerateContentConfig +import com.google.adk.kt.types.Part +import com.google.adk.kt.types.Role +import com.google.genai.types.CachedContent +import com.google.genai.types.CreateCachedContentConfig +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertNull +import kotlin.time.Clock + +class GeminiContextCacheManagerTest { + + /** A [GeminiContextCacheManager.CacheClient] that records calls and returns a fixed cache. */ + private class FakeCacheClient(private val createdName: String = "cache/new") : + GeminiContextCacheManager.CacheClient { + var createCount = 0 + var deleteCount = 0 + var lastDeletedName: String? = null + + override fun create(model: String, config: CreateCachedContentConfig): CachedContent { + createCount++ + return CachedContent.builder().name(createdName).build() + } + + override fun delete(name: String) { + deleteCount++ + lastDeletedName = name + } + } + + private val cacheConfig = ContextCacheConfig() + + private fun textContent(role: String, text: String) = + Content(role = role, parts = listOf(Part(text = text))) + + // A request whose cacheable prefix is the leading model content (count = 1), leaving the trailing + // user content to send. + private fun baseRequest(cacheMetadata: CacheMetadata? = null, tokenCount: Int? = null) = + LlmRequest( + contents = listOf(textContent(Role.MODEL, "cached"), textContent(Role.USER, "latest")), + config = GenerateContentConfig(systemInstruction = textContent(Role.MODEL, "be helpful")), + cacheConfig = cacheConfig, + cacheMetadata = cacheMetadata, + cacheableContentsTokenCount = tokenCount, + ) + + private fun fingerprintFor(manager: GeminiContextCacheManager, request: LlmRequest): String = + manager.handleContextCaching(request.copy(cacheMetadata = null)).cacheMetadata!!.fingerprint + + @Test + fun handleContextCaching_noExistingMetadata_returnsFingerprintOnlyAndDoesNotCreate() { + val fake = FakeCacheClient() + val manager = GeminiContextCacheManager("gemini-2.0-flash", fake) + val request = baseRequest() + + val result = manager.handleContextCaching(request) + + assertNull(result.cacheMetadata?.cacheName) + assertEquals(1, result.cacheMetadata?.contentsCount) + // Request is left unchanged (no cache applied). + assertEquals(request.contents, result.request.contents) + assertNull(result.request.config.cachedContent) + assertEquals(0, fake.createCount) + } + + @Test + fun handleContextCaching_validActiveCache_appliesCacheToRequest() { + val fake = FakeCacheClient() + val manager = GeminiContextCacheManager("gemini-2.0-flash", fake) + val request = baseRequest() + val activeMetadata = + CacheMetadata( + fingerprint = fingerprintFor(manager, request), + contentsCount = 1, + cacheName = "cache/existing", + expireTime = Clock.System.now().toEpochMilliseconds() + 100_000, + invocationsUsed = 1, + ) + + val result = manager.handleContextCaching(request.copy(cacheMetadata = activeMetadata)) + + assertEquals("cache/existing", result.request.config.cachedContent) + assertNull(result.request.config.systemInstruction) + assertEquals(1, result.request.contents.size) + assertEquals("latest", result.request.contents[0].parts[0].text) + assertEquals(activeMetadata, result.cacheMetadata) + assertEquals(0, fake.createCount) + assertEquals(0, fake.deleteCount) + } + + @Test + fun handleContextCaching_expiredCacheSameFingerprintEnoughTokens_recreatesCache() { + val fake = FakeCacheClient(createdName = "cache/recreated") + val manager = GeminiContextCacheManager("gemini-2.0-flash", fake) + val request = baseRequest(tokenCount = 5000) + val expiredMetadata = + CacheMetadata( + fingerprint = fingerprintFor(manager, request), + contentsCount = 1, + cacheName = "cache/old", + expireTime = Clock.System.now().toEpochMilliseconds() - 1_000, + invocationsUsed = 1, + ) + + val result = + manager.handleContextCaching( + request.copy(cacheMetadata = expiredMetadata, cacheableContentsTokenCount = 5000) + ) + + assertEquals(1, fake.deleteCount) + assertEquals("cache/old", fake.lastDeletedName) + assertEquals(1, fake.createCount) + assertEquals("cache/recreated", result.request.config.cachedContent) + assertEquals("cache/recreated", result.cacheMetadata?.cacheName) + assertEquals(1, result.cacheMetadata?.invocationsUsed) + } + + @Test + fun handleContextCaching_expiredCacheTooFewTokens_returnsFingerprintOnly() { + val fake = FakeCacheClient() + val manager = GeminiContextCacheManager("gemini-2.0-flash", fake) + val request = baseRequest() + val expiredMetadata = + CacheMetadata( + fingerprint = fingerprintFor(manager, request), + contentsCount = 1, + cacheName = "cache/old", + expireTime = Clock.System.now().toEpochMilliseconds() - 1_000, + invocationsUsed = 1, + ) + + // No token count available, so the cache cannot be recreated. + val result = manager.handleContextCaching(request.copy(cacheMetadata = expiredMetadata)) + + assertEquals(1, fake.deleteCount) + assertEquals(0, fake.createCount) + assertNull(result.cacheMetadata?.cacheName) + assertEquals(1, result.cacheMetadata?.contentsCount) + } + + @Test + fun populateCacheMetadataInResponse_attachesMetadata() { + val manager = GeminiContextCacheManager("gemini-2.0-flash", FakeCacheClient()) + val response = LlmResponse(content = textContent(Role.MODEL, "ok")) + val metadata = CacheMetadata(fingerprint = "fp", contentsCount = 0) + + val result = manager.populateCacheMetadataInResponse(response, metadata) + + assertEquals(metadata, result.cacheMetadata) + } +}