diff --git a/core/src/commonJvmAndroidMain/kotlin/com/google/adk/kt/crypto/Sha256.kt b/core/src/commonJvmAndroidMain/kotlin/com/google/adk/kt/crypto/Sha256.kt new file mode 100644 index 00000000..22b8718d --- /dev/null +++ b/core/src/commonJvmAndroidMain/kotlin/com/google/adk/kt/crypto/Sha256.kt @@ -0,0 +1,24 @@ +/* + * 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.crypto + +import java.security.MessageDigest + +internal actual fun sha256Hex(input: String): String = + MessageDigest.getInstance("SHA-256").digest(input.encodeToByteArray()).joinToString("") { + (it.toInt() and 0xff).toString(16).padStart(2, '0') + } 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/crypto/Sha256.kt b/core/src/commonMain/kotlin/com/google/adk/kt/crypto/Sha256.kt new file mode 100644 index 00000000..96b14b63 --- /dev/null +++ b/core/src/commonMain/kotlin/com/google/adk/kt/crypto/Sha256.kt @@ -0,0 +1,20 @@ +/* + * 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.crypto + +/** Returns the lowercase hex-encoded SHA-256 digest of [input]'s UTF-8 bytes. */ +internal expect fun sha256Hex(input: String): String diff --git a/core/src/commonMain/kotlin/com/google/adk/kt/models/Gemini.kt b/core/src/commonMain/kotlin/com/google/adk/kt/models/Gemini.kt index 877747de..3b04d9c2 100644 --- a/core/src/commonMain/kotlin/com/google/adk/kt/models/Gemini.kt +++ b/core/src/commonMain/kotlin/com/google/adk/kt/models/Gemini.kt @@ -132,10 +132,21 @@ class Gemini( override fun generateContent(request: LlmRequest, stream: Boolean): Flow = flow { val preparedRequest = request.prepareGenerateContentRequest(!client.enterprise) - val config = preparedRequest.config.toGenaiSdk() - val contents = preparedRequest.contents.map { it.toGenaiSdk() } - logger.debug { "LLM Request:\n${Json.toJsonString(buildLoggingRequestMap(preparedRequest))}" } + // 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, GenaiCacheClient(client.caches)) + } + val cacheResult = cacheManager?.handleContextCaching(preparedRequest) + val finalRequest = cacheResult?.request ?: preparedRequest + val cacheMetadata = cacheResult?.cacheMetadata + + val config = finalRequest.config.toGenaiSdk() + val contents = finalRequest.contents.map { it.toGenaiSdk() } + + logger.debug { "LLM Request:\n${Json.toJsonString(buildLoggingRequestMap(finalRequest))}" } if (stream) { val aggregator = StreamingResponseAggregator() @@ -148,8 +159,8 @@ class Gemini( emit(aggregator.processResponse(response.fromGenaiSdk())) } - // After stream loop ends, emit final aggregated response - aggregator.aggregate()?.let { emit(it) } + // After stream loop ends, emit final aggregated response with any cache metadata attached + aggregator.aggregate()?.let { emit(it.copy(cacheMetadata = cacheMetadata)) } } else { val response = models.generateContent(name, contents, config) logger.debug { @@ -157,7 +168,7 @@ class Gemini( "finishReason=${response.candidates?.firstOrNull()?.finishReason}" } val llmResponse = LlmResponse.from(response.fromGenaiSdk()) - emit(llmResponse) + emit(llmResponse.copy(cacheMetadata = cacheMetadata)) } } diff --git a/core/src/commonMain/kotlin/com/google/adk/kt/models/GeminiContextCacheManager.kt b/core/src/commonMain/kotlin/com/google/adk/kt/models/GeminiContextCacheManager.kt new file mode 100644 index 00000000..6c22d905 --- /dev/null +++ b/core/src/commonMain/kotlin/com/google/adk/kt/models/GeminiContextCacheManager.kt @@ -0,0 +1,406 @@ +/* + * 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.annotations.FrameworkInternalApi +import com.google.adk.kt.crypto.sha256Hex +import com.google.adk.kt.logging.LoggerFactory +import com.google.adk.kt.serialization.adkJson +import com.google.adk.kt.types.Content +import com.google.adk.kt.types.Role +import com.google.adk.kt.types.Tool +import com.google.adk.kt.types.ToolConfig +import kotlin.time.Clock +import kotlin.time.Duration +import kotlinx.serialization.json.JsonArray +import kotlinx.serialization.json.JsonElement +import kotlinx.serialization.json.JsonObject +import kotlinx.serialization.json.JsonPrimitive +import kotlinx.serialization.json.encodeToJsonElement + +/** + * 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. + * + * The cache backend is delegated to a [CacheClient] interface so the manager stays decoupled from + * the GenAI SDK and can be faked in tests: the SDK's cache types are final with internal + * constructors and cannot be faked directly. The Gemini model wires in the SDK-backed + * `GenaiCacheClient`. + * + * @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. + */ + suspend 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. Only a real + // cache has a backend resource to delete; a fingerprint-only record (cacheName == null) has + // nothing to delete. + val staleCacheName = existing.cacheName + if (staleCacheName != null) { + logger.debug { "Cache is invalid, cleaning up: $staleCacheName" } + cleanupCache(staleCacheName) + } + + val cacheContentsCount = existing.contentsCount + val currentFingerprint = generateFingerprint(request, cacheContentsCount) + + if (currentFingerprint == existing.fingerprint) { + // The previously cached prefix is unchanged (fingerprints match) but the cache expired. + // Recreate it, growing the prefix to the current cacheable boundary (newly settled history) + // so it keeps up as the conversation grows, matching Python's + // max(previousCount, currentCacheableCount). + val newCount = maxOf(cacheContentsCount, findCountOfContentsToCache(request.contents)) + val newFingerprint = generateFingerprint(request, newCount) + val created = createNewCacheWithContents(request, newCount) + if (created != null) { + val applied = applyCacheToRequest(request, created.cacheName!!, newCount) + return CacheResult(applied, created) + } + // Creation failed (e.g. below the token minimum): preserve the grown prefix fingerprint so + // the fingerprint chain stays stable for subsequent calls. + return CacheResult( + request, + CacheMetadata(fingerprint = newFingerprint, contentsCount = newCount), + ) + } + + // 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), + ) + } + + /** + * 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. + * + * Callers run `ensureModelResponse` (via `prepareGenerateContentRequest`) before caching, so the + * last content is always a user turn; the `contents.size` fallback for a non-user or empty tail + * is therefore unreachable in the real flow. + */ + 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 cacheIntervals = request.cacheConfig?.cacheIntervals ?: return false + if (metadata.invocationsUsed!! > cacheIntervals) { + logger.info { + "Cache exceeded cache intervals: ${metadata.cacheName} " + + "(${metadata.invocationsUsed} > $cacheIntervals)" + } + 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: model, system instruction, + * tools, tool config, and the first [cacheContentsCount] contents. Serializes them into a single + * canonical JSON object with every object's keys recursively sorted, so the fingerprint depends + * only on content, not on map insertion or property order, and has no field-boundary ambiguity. + * The model is included because explicit caches are model-specific. + */ + @OptIn(FrameworkInternalApi::class) + private fun generateFingerprint(request: LlmRequest, cacheContentsCount: Int): String { + val fields = + buildMap { + this["model"] = JsonPrimitive(modelName) + request.config.systemInstruction?.let { + this["system_instruction"] = adkJson.encodeToJsonElement(it) + } + request.config.tools?.let { this["tools"] = adkJson.encodeToJsonElement(it) } + request.config.toolConfig?.let { this["tool_config"] = adkJson.encodeToJsonElement(it) } + if (cacheContentsCount > 0 && request.contents.isNotEmpty()) { + val count = minOf(cacheContentsCount, request.contents.size) + this["cached_contents"] = adkJson.encodeToJsonElement(request.contents.take(count)) + } + } + return sha256Hex(JsonObject(fields).sortedKeys().toString()).take(FINGERPRINT_LENGTH) + } + + /** + * Returns a copy with every nested [JsonObject]'s keys sorted alphabetically (array element order + * is preserved). Objects with the same entries then serialize identically regardless of the order + * their keys were inserted, so the fingerprint is stable across equivalent requests. + */ + private fun JsonElement.sortedKeys(): JsonElement = + when (this) { + is JsonObject -> + JsonObject(entries.sortedBy { it.key }.associate { (k, v) -> k to v.sortedKeys() }) + is JsonArray -> JsonArray(map { it.sortedKeys() }) + else -> this + } + + /** Creates a new cache when the previous request was large enough, otherwise returns `null`. */ + private suspend 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 + } + // `tokenCount` covers the whole previous prompt, but the cache stores only the prefix + // (system instruction + tools + first N contents). Gate the model's documented floor on the + // estimated prefix size so we never send a sub-minimum payload that Gemini would reject with + // 400. Opaque models have no known floor, so the server stays authoritative. + val cacheablePrefixTokens = estimateCacheablePrefixTokens(request, cacheContentsCount) + val minimumTokens = minimumCacheTokens(modelName) + if (minimumTokens != null && cacheablePrefixTokens < minimumTokens) { + logger.info { + "Cacheable prefix below Gemini minimum cache size " + + "($cacheablePrefixTokens < $minimumTokens tokens)." + } + return null + } + return try { + createGeminiCache(request, cacheContentsCount) + } catch (e: Exception) { + logger.warn { "Failed to create cache: ${e.message}" } + null + } + } + + /** + * The explicit-cache token floor for a named Gemini model family, or `null` for opaque + * tuned-model / endpoint IDs where the server remains authoritative. + */ + private fun minimumCacheTokens(model: String): Int? { + val bareName = model.substringAfterLast('/') + return when { + bareName.startsWith("gemini-2.5-") -> GEMINI_2_5_MIN_CACHE_TOKENS + bareName.startsWith("gemini-3") -> GEMINI_3_MIN_CACHE_TOKENS + else -> null + } + } + + /** Creates the cache via the [cacheClient] and returns its metadata. */ + private suspend fun createGeminiCache( + request: LlmRequest, + cacheContentsCount: Int, + ): CacheMetadata { + val cacheConfig = requireNotNull(request.cacheConfig) { "cacheConfig must be set." } + val displayName = "adk-cache-${Clock.System.now().epochSeconds}-${cacheContentsCount}contents" + + val cacheName = + cacheClient.create( + CacheCreateRequest( + model = modelName, + contents = request.contents.take(cacheContentsCount).takeIf { it.isNotEmpty() }, + systemInstruction = request.config.systemInstruction, + tools = request.config.tools, + toolConfig = request.config.toolConfig, + ttl = cacheConfig.ttl, + displayName = displayName, + ) + ) + val createdAt = Clock.System.now().toEpochMilliseconds() + 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 suspend 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, + toolConfig = null, + cachedContent = cacheName, + ), + contents = request.contents.drop(cacheContentsCount), + ) + + /** The inputs needed to create a Gemini cache, expressed in ADK common types. */ + data class CacheCreateRequest( + val model: String, + // `null` when the cache holds no contents (system instruction + tools only), matching the API's + // distinction between an absent and an empty contents field. + val contents: List?, + val systemInstruction: Content?, + val tools: List?, + val toolConfig: ToolConfig?, + val ttl: Duration, + val displayName: String, + ) + + /** + * Abstraction over the cache backend so the manager can be faked in tests; the SDK's cache types + * are final with internal constructors and cannot be faked directly. The SDK-backed + * implementation is [GenaiCacheClient]. + */ + interface CacheClient { + /** Creates a cache and returns its resource name. */ + suspend fun create(request: CacheCreateRequest): String + + suspend fun delete(name: String) + } + + /** + * Estimates the token count of the prefix that will actually be cached (system instruction, + * tools, and the first [cacheContentsCount] contents). The only accurate count available is + * [LlmRequest.cacheableContentsTokenCount], which covers the whole previous prompt, so this + * scales it by the prefix's estimated character share of the request. + */ + @OptIn(FrameworkInternalApi::class) + private fun estimateCacheablePrefixTokens(request: LlmRequest, cacheContentsCount: Int): Int { + val fullTokens = request.cacheableContentsTokenCount ?: return 0 + if (fullTokens == 0) return 0 + val fullEstimate = estimateRequestTokens(request, null) + // No text to estimate from (e.g. binary-only parts): trust the accurate full count. + if (fullEstimate <= 0) return fullTokens + val prefixEstimate = estimateRequestTokens(request, cacheContentsCount) + val ratio = minOf(1.0, prefixEstimate.toDouble() / fullEstimate) + return (fullTokens * ratio).toInt() + } + + /** + * Rough character-based token estimate (~4 chars/token) for the request, or its cacheable prefix + * when [cacheContentsCount] is non-null. Always counts system instruction and tools. + */ + @OptIn(FrameworkInternalApi::class) + private fun estimateRequestTokens(request: LlmRequest, cacheContentsCount: Int?): Int { + var totalChars = 0 + request.config.systemInstruction?.parts?.forEach { totalChars += it.text?.length ?: 0 } + request.config.tools?.let { tools -> + for (tool in tools) { + totalChars += adkJson.encodeToString(Tool.serializer(), tool).length + } + } + val contents = + if (cacheContentsCount != null) request.contents.take(cacheContentsCount) + else request.contents + for (content in contents) { + for (part in content.parts) { + totalChars += part.text?.length ?: 0 + } + } + return totalChars / CHARS_PER_TOKEN + } + + private companion object { + private val logger = LoggerFactory.getLogger(GeminiContextCacheManager::class) + + // Named Gemini model families have documented explicit-cache token floors; opaque tuned-model + // and endpoint IDs leave the floor to the server. + private const val GEMINI_2_5_MIN_CACHE_TOKENS = 2048 + private const val GEMINI_3_MIN_CACHE_TOKENS = 4096 + + private const val FINGERPRINT_LENGTH = 16 + + // Rough heuristic for estimating tokens from character length. + private const val CHARS_PER_TOKEN = 4 + } +} diff --git a/core/src/commonMain/kotlin/com/google/adk/kt/models/GenaiCacheClient.kt b/core/src/commonMain/kotlin/com/google/adk/kt/models/GenaiCacheClient.kt new file mode 100644 index 00000000..197b78df --- /dev/null +++ b/core/src/commonMain/kotlin/com/google/adk/kt/models/GenaiCacheClient.kt @@ -0,0 +1,48 @@ +/* + * 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.types.toGenaiSdk +import com.google.genai.kotlin.Caches +import com.google.genai.kotlin.types.CreateCachedContentConfig + +/** + * [GeminiContextCacheManager.CacheClient] backed by the GenAI SDK's `client.caches`. + * + * This is the single place that touches the GenAI SDK cache API, translating ADK common types to + * and from the SDK so the manager itself stays decoupled from the SDK and can be faked in tests. + */ +internal class GenaiCacheClient(private val caches: Caches) : + GeminiContextCacheManager.CacheClient { + + override suspend fun create(request: GeminiContextCacheManager.CacheCreateRequest): String { + val config = + CreateCachedContentConfig( + contents = request.contents?.map { it.toGenaiSdk() }, + systemInstruction = request.systemInstruction?.toGenaiSdk(), + tools = request.tools?.map { it.toGenaiSdk() }, + toolConfig = request.toolConfig?.toGenaiSdk(), + ttl = request.ttl, + displayName = request.displayName, + ) + val cachedContent = caches.create(request.model, config) + return cachedContent.name ?: throw IllegalStateException("Created cache has no resource name.") + } + + override suspend fun delete(name: String) { + val unused = caches.delete(name) + } +} 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 5adc4e1e..1c7b0d5b 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 @@ -546,7 +546,7 @@ abstract class AbstractRunner : Runner { pluginManager = pluginManager, resumabilityConfig = resumabilityConfig, eventsCompactionConfig = app?.eventsCompactionConfig, - contextCacheConfig = null, + contextCacheConfig = app?.contextCacheConfig, ) .let { // Run callbacks and append user message to session @@ -603,7 +603,7 @@ abstract class AbstractRunner : Runner { pluginManager = pluginManager, resumabilityConfig = resumabilityConfig, eventsCompactionConfig = app?.eventsCompactionConfig, - contextCacheConfig = null, + contextCacheConfig = app?.contextCacheConfig, ) val currentContext = 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 d37cdc2b..57327514 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 @@ -14,8 +14,11 @@ * limitations under the License. */ +@file:OptIn(com.google.adk.kt.annotations.ExperimentalContextCachingFeature::class) + 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.plugins.Plugin import com.google.adk.kt.summarizer.EventsCompactionConfig @@ -54,6 +57,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(cacheIntervals = 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/GeminiContextCacheManagerTest.kt b/core/src/commonTest/kotlin/com/google/adk/kt/models/GeminiContextCacheManagerTest.kt new file mode 100644 index 00000000..ed41c24a --- /dev/null +++ b/core/src/commonTest/kotlin/com/google/adk/kt/models/GeminiContextCacheManagerTest.kt @@ -0,0 +1,404 @@ +/* + * 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.models + +import com.google.adk.kt.agents.ContextCacheConfig +import com.google.adk.kt.types.Content +import com.google.adk.kt.types.FunctionCall +import com.google.adk.kt.types.GenerateContentConfig +import com.google.adk.kt.types.Part +import com.google.adk.kt.types.Role +import com.google.adk.kt.types.ToolConfig +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertNull +import kotlin.time.Clock +import kotlinx.coroutines.test.runTest + +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 + var lastCreateRequest: GeminiContextCacheManager.CacheCreateRequest? = null + + override suspend fun create(request: GeminiContextCacheManager.CacheCreateRequest): String { + createCount++ + lastCreateRequest = request + return createdName + } + + override suspend 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 suspend fun fingerprintFor( + manager: GeminiContextCacheManager, + request: LlmRequest, + ): String = + manager.handleContextCaching(request.copy(cacheMetadata = null)).cacheMetadata!!.fingerprint + + // A request whose cacheable prefix (a large system instruction, no cached contents) estimates to + // ~3000 tokens: above Gemini 2.5's 2048 floor but below Gemini 3's 4096 floor. + private fun floorRequest() = + LlmRequest( + contents = listOf(textContent(Role.USER, "hi")), + config = + GenerateContentConfig(systemInstruction = textContent(Role.MODEL, "x".repeat(12_000))), + cacheConfig = cacheConfig, + cacheableContentsTokenCount = 3000, + ) + + @Test + fun handleContextCaching_noExistingMetadata_returnsFingerprintOnlyAndDoesNotCreate() = runTest { + 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() = runTest { + 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() = runTest { + val fake = FakeCacheClient(createdName = "cache/recreated") + val manager = GeminiContextCacheManager("gemini-2.0-flash", fake) + // Large enough that the estimated cacheable prefix still clears Gemini's 4096-token floor. + val request = baseRequest(tokenCount = 8000) + 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 = 8000) + ) + + assertEquals(1, fake.deleteCount) + assertEquals("cache/old", fake.lastDeletedName) + assertEquals(1, fake.createCount) + // The non-empty cached prefix (the leading model content) is passed through to create. + assertEquals(1, fake.lastCreateRequest?.contents?.size) + assertEquals("cache/recreated", result.request.config.cachedContent) + assertEquals("cache/recreated", result.cacheMetadata?.cacheName) + assertEquals(1, result.cacheMetadata?.invocationsUsed) + } + + @Test + fun handleContextCaching_expiredCacheConversationGrew_recreatesWithGrownPrefix() = runTest { + val fake = FakeCacheClient(createdName = "cache/grown") + val manager = GeminiContextCacheManager("gemini-2.0-flash", fake) + // The count-1 prefix ([model "cached"]) is unchanged, so the old count-1 fingerprint still + // matches; but the conversation grew (findCount is now 2), so recreation must cache the larger + // prefix (matches Python's max(previousCount, currentCacheableCount)). + val shortRequest = baseRequest() + val fingerprintAtCountOne = fingerprintFor(manager, shortRequest) + val grownRequest = + shortRequest.copy( + contents = + listOf( + textContent(Role.MODEL, "cached"), + textContent(Role.MODEL, "more"), + textContent(Role.USER, "latest"), + ) + ) + val expiredMetadata = + CacheMetadata( + fingerprint = fingerprintAtCountOne, + contentsCount = 1, + cacheName = "cache/old", + expireTime = Clock.System.now().toEpochMilliseconds() - 1_000, + invocationsUsed = 1, + ) + + val result = + manager.handleContextCaching( + grownRequest.copy(cacheMetadata = expiredMetadata, cacheableContentsTokenCount = 8000) + ) + + assertEquals(1, fake.createCount) + // The recreated cache grew from 1 to 2 contents. + assertEquals(2, fake.lastCreateRequest?.contents?.size) + assertEquals(2, result.cacheMetadata?.contentsCount) + } + + @Test + fun handleContextCaching_expiredCacheTooFewTokens_returnsFingerprintOnly() = runTest { + 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 handleContextCaching_differentModel_doesNotReuseCache() = runTest { + val request = baseRequest() + // Fingerprint captured under one model must not validate an active cache under another, because + // explicit caches are model-specific (the model is part of the fingerprint). + val flashFingerprint = + fingerprintFor(GeminiContextCacheManager("gemini-2.0-flash", FakeCacheClient()), request) + val activeMetadata = + CacheMetadata( + fingerprint = flashFingerprint, + contentsCount = 1, + cacheName = "cache/flash", + expireTime = Clock.System.now().toEpochMilliseconds() + 100_000, + invocationsUsed = 1, + ) + + val proFake = FakeCacheClient() + val proManager = GeminiContextCacheManager("gemini-1.5-pro", proFake) + val result = proManager.handleContextCaching(request.copy(cacheMetadata = activeMetadata)) + + assertNull(result.request.config.cachedContent) + assertEquals(1, proFake.deleteCount) + assertEquals("cache/flash", proFake.lastDeletedName) + assertNull(result.cacheMetadata?.cacheName) + } + + @Test + fun handleContextCaching_validCache_stripsToolConfig() = runTest { + val fake = FakeCacheClient() + val manager = GeminiContextCacheManager("gemini-2.0-flash", fake) + val request = baseRequest().let { it.copy(config = it.config.copy(toolConfig = ToolConfig())) } + 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.toolConfig) + } + + @Test + fun handleContextCaching_prefixBelowMinimumEvenWhenFullClears_doesNotCreate() = runTest { + val fake = FakeCacheClient() + val manager = GeminiContextCacheManager("gemini-2.5-flash", fake) + // Tiny cached prefix but a huge trailing (uncached) user turn: the full previous prompt clears + // Gemini 2.5's 2048-token floor, but the estimated cacheable prefix does not, so no cache is + // created. + val request = + LlmRequest( + contents = + listOf(textContent(Role.MODEL, "cached"), textContent(Role.USER, "y".repeat(100_000))), + config = GenerateContentConfig(systemInstruction = textContent(Role.MODEL, "be helpful")), + cacheConfig = cacheConfig, + cacheableContentsTokenCount = 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(0, fake.createCount) + assertNull(result.cacheMetadata?.cacheName) + } + + @Test + fun handleContextCaching_contentsCount_cachesBeforeLastUserBatch() = runTest { + val manager = GeminiContextCacheManager("gemini-2.0-flash", FakeCacheClient()) + + suspend fun countFor(vararg roles: String): Int { + val contents = roles.mapIndexed { i, role -> textContent(role, "m$i") } + val request = LlmRequest(contents = contents, cacheConfig = cacheConfig) + return manager.handleContextCaching(request).cacheMetadata!!.contentsCount + } + + // Empty contents: nothing to cache. + assertEquals(0, countFor()) + // A single user turn has nothing before it to cache. + assertEquals(0, countFor(Role.USER)) + // Cache everything before the trailing contiguous user batch. + assertEquals(2, countFor(Role.USER, Role.MODEL, Role.USER)) + assertEquals(4, countFor(Role.USER, Role.MODEL, Role.USER, Role.MODEL, Role.USER)) + // A model turn last means no trailing user batch: the degenerate "cache everything" count. + assertEquals(6, countFor(Role.USER, Role.MODEL, Role.USER, Role.MODEL, Role.USER, Role.MODEL)) + } + + @Test + fun handleContextCaching_equivalentArgMappingOrder_produceSameFingerprint() = runTest { + val manager = GeminiContextCacheManager("gemini-2.0-flash", FakeCacheClient()) + + fun requestWithArgs(args: Map) = + LlmRequest( + contents = + listOf( + Content( + role = Role.MODEL, + parts = listOf(Part(functionCall = FunctionCall(name = "lookup", args = args))), + ), + textContent(Role.USER, "latest"), + ), + config = GenerateContentConfig(systemInstruction = textContent(Role.MODEL, "be helpful")), + cacheConfig = cacheConfig, + ) + + // Same args in different key insertion order: the canonical fingerprint must ignore order. + val first = fingerprintFor(manager, requestWithArgs(mapOf("first" to 1, "second" to 2))) + val second = fingerprintFor(manager, requestWithArgs(mapOf("second" to 2, "first" to 1))) + + assertEquals(first, second) + } + + @Test + fun handleContextCaching_gemini25_createsCacheAbove2048Floor() = runTest { + val fake = FakeCacheClient(createdName = "cache/gemini25") + val manager = GeminiContextCacheManager("gemini-2.5-flash", fake) + val request = floorRequest() + val metadata = CacheMetadata(fingerprint = fingerprintFor(manager, request), contentsCount = 0) + + val result = manager.handleContextCaching(request.copy(cacheMetadata = metadata)) + + // ~3000 estimated prefix tokens clears Gemini 2.5's 2048-token floor. + assertEquals(1, fake.createCount) + assertEquals("cache/gemini25", result.cacheMetadata?.cacheName) + } + + @Test + fun handleContextCaching_gemini3_skipsCacheBelow4096Floor() = runTest { + val fake = FakeCacheClient() + val manager = GeminiContextCacheManager("gemini-3.1-pro-preview", fake) + val request = floorRequest() + val metadata = CacheMetadata(fingerprint = fingerprintFor(manager, request), contentsCount = 0) + + val result = manager.handleContextCaching(request.copy(cacheMetadata = metadata)) + + // ~3000 estimated prefix tokens is below Gemini 3's 4096-token floor. + assertEquals(0, fake.createCount) + assertNull(result.cacheMetadata?.cacheName) + } + + @Test + fun handleContextCaching_opaqueModel_leavesTokenFloorToServer() = runTest { + val fake = FakeCacheClient(createdName = "cache/tuned") + val manager = + GeminiContextCacheManager("projects/p/locations/us-central1/endpoints/tuned-model", fake) + val request = floorRequest() + val metadata = CacheMetadata(fingerprint = fingerprintFor(manager, request), contentsCount = 0) + + val result = manager.handleContextCaching(request.copy(cacheMetadata = metadata)) + + // Opaque model IDs have no client-side floor, so creation is attempted (the server decides). + assertEquals(1, fake.createCount) + assertEquals("cache/tuned", result.cacheMetadata?.cacheName) + } + + @Test + fun handleContextCaching_createWithEmptyPrefix_passesNullContents() = runTest { + val fake = FakeCacheClient(createdName = "cache/sysonly") + val manager = GeminiContextCacheManager("gemini-2.5-flash", fake) + // All-user contents: the cached prefix is 0 contents (system instruction + tools only). + val request = floorRequest() + val metadata = CacheMetadata(fingerprint = fingerprintFor(manager, request), contentsCount = 0) + + val result = manager.handleContextCaching(request.copy(cacheMetadata = metadata)) + + assertEquals(1, fake.createCount) + // An empty prefix is passed as null (an absent contents field), not an empty list. + assertNull(fake.lastCreateRequest?.contents) + assertEquals("cache/sysonly", result.cacheMetadata?.cacheName) + } +} 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 2288f21c..3e7dfe68 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 @@ -936,6 +936,31 @@ class AbstractRunnerTest { // ----- Context cache config propagation ----- + @Test + fun runAsync_appWithContextCacheConfig_propagatesToInvocationContext() = runTest { + val cacheConfig = ContextCacheConfig(cacheIntervals = 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.