diff --git a/tokt/src/main/kotlin/com/google/adk/tokt/codecs/CacheMetadataCodec.kt b/tokt/src/main/kotlin/com/google/adk/tokt/codecs/CacheMetadataCodec.kt new file mode 100644 index 000000000..150fabb14 --- /dev/null +++ b/tokt/src/main/kotlin/com/google/adk/tokt/codecs/CacheMetadataCodec.kt @@ -0,0 +1,52 @@ +/* + * 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.tokt.codecs + +import com.google.adk.kt.models.CacheMetadata as KtCacheMetadata +import com.google.adk.models.CacheMetadata as JavaCacheMetadata +import java.time.Instant +import kotlin.jvm.optionals.getOrNull + +/** + * Converts context cache metadata between ADK Kotlin and ADK Java. Cache times are epoch + * milliseconds on the Kotlin side and [Instant]s on the Java side. Shared by [EventCodec] and + * [LlmResponseCodec]. + */ +internal object CacheMetadataCodec { + + /** Returns the Java [JavaCacheMetadata] view of the Kotlin [metadata]. */ + fun toJava(metadata: KtCacheMetadata): JavaCacheMetadata = + JavaCacheMetadata.builder() + .fingerprint(metadata.fingerprint) + .contentsCount(metadata.contentsCount) + .cacheName(metadata.cacheName) + .expireTime(metadata.expireTime?.let { Instant.ofEpochMilli(it) }) + .invocationsUsed(metadata.invocationsUsed) + .createdAt(metadata.createdAt?.let { Instant.ofEpochMilli(it) }) + .build() + + /** Returns the Kotlin [KtCacheMetadata] view of the Java [metadata]. */ + fun fromJava(metadata: JavaCacheMetadata): KtCacheMetadata = + KtCacheMetadata( + fingerprint = metadata.fingerprint(), + contentsCount = metadata.contentsCount(), + cacheName = metadata.cacheName().getOrNull(), + expireTime = metadata.expireTime().getOrNull()?.toEpochMilli(), + invocationsUsed = metadata.invocationsUsed().getOrNull(), + createdAt = metadata.createdAt().getOrNull()?.toEpochMilli(), + ) +} diff --git a/tokt/src/main/kotlin/com/google/adk/tokt/codecs/EventCodec.kt b/tokt/src/main/kotlin/com/google/adk/tokt/codecs/EventCodec.kt index f05f264ab..59fc574cb 100644 --- a/tokt/src/main/kotlin/com/google/adk/tokt/codecs/EventCodec.kt +++ b/tokt/src/main/kotlin/com/google/adk/tokt/codecs/EventCodec.kt @@ -33,9 +33,10 @@ import kotlin.jvm.optionals.getOrNull * skipSummarization, endOfAgent - the requested tool confirmations, the context-compaction summary, * the resumable-run `agentState`, the rewind marker `rewindBeforeInvocationId`, and the workflow * `route`), the streaming/error signals (`partial`, `turnComplete`, `interrupted`, `errorMessage`, - * `errorCode`, `finishReason`), the token `usageMetadata` ([UsageMetadataCodec]) and - * `groundingMetadata` ([GroundingMetadataCodec]), a workflow node's `output` and `nodeInfo` - * ([WorkflowCodec]), long-running tool ids, and the timestamp. + * `errorCode`, `finishReason`), the token `usageMetadata` ([UsageMetadataCodec]), + * `groundingMetadata` ([GroundingMetadataCodec]) and context `cacheMetadata` + * ([CacheMetadataCodec]), a workflow node's `output` and `nodeInfo` ([WorkflowCodec]), long-running + * tool ids, and the timestamp. */ internal object EventCodec { @@ -61,6 +62,7 @@ internal object EventCodec { event.usageMetadata?.let { builder.usageMetadata(UsageMetadataCodec.toJava(it)) } event.avgLogProbs?.let { builder.avgLogprobs(it) } event.groundingMetadata?.let { builder.groundingMetadata(GroundingMetadataCodec.toJava(it)) } + event.cacheMetadata?.let { builder.cacheMetadata(CacheMetadataCodec.toJava(it)) } event.modelVersion?.let { builder.modelVersion(it) } event.customMetadata?.let { builder.customMetadata(CustomMetadataCodec.toJava(it)) } event.output?.let { builder.output(WorkflowCodec.outputToJava(it)) } @@ -111,6 +113,7 @@ internal object EventCodec { avgLogProbs = event.avgLogprobs().getOrNull(), groundingMetadata = event.groundingMetadata().getOrNull()?.let { GroundingMetadataCodec.fromJava(it) }, + cacheMetadata = event.cacheMetadata().getOrNull()?.let { CacheMetadataCodec.fromJava(it) }, modelVersion = event.modelVersion().getOrNull(), customMetadata = event.customMetadata().getOrNull()?.let { CustomMetadataCodec.fromJava(it) }, output = event.output().getOrNull()?.let { WorkflowCodec.outputFromJava(it) }, diff --git a/tokt/src/main/kotlin/com/google/adk/tokt/codecs/LlmResponseCodec.kt b/tokt/src/main/kotlin/com/google/adk/tokt/codecs/LlmResponseCodec.kt index c5544a737..22ac80309 100644 --- a/tokt/src/main/kotlin/com/google/adk/tokt/codecs/LlmResponseCodec.kt +++ b/tokt/src/main/kotlin/com/google/adk/tokt/codecs/LlmResponseCodec.kt @@ -28,8 +28,8 @@ import kotlin.jvm.optionals.getOrNull * a bridged Java plugin's `afterModel`. * * Carries content, error message, streaming flags, model version, finish reason, usage metadata, - * grounding metadata, average log-probs, error code, and custom metadata. Citation metadata and - * logprobs results have no ADK Java facade field, so they are not mapped. + * grounding and context cache metadata, average log-probs, error code, and custom metadata. + * Citation metadata and logprobs results have no ADK Java facade field, so they are not mapped. */ internal object LlmResponseCodec { @@ -45,6 +45,7 @@ internal object LlmResponseCodec { usageMetadata = response.usageMetadata().getOrNull()?.let { UsageMetadataCodec.fromJava(it) }, groundingMetadata = response.groundingMetadata().getOrNull()?.let { GroundingMetadataCodec.fromJava(it) }, + cacheMetadata = response.cacheMetadata().getOrNull()?.let { CacheMetadataCodec.fromJava(it) }, avgLogprobs = response.avgLogprobs().getOrNull(), // toString() preserves the raw value; knownEnum().name would collapse an unrecognized code to // FINISH_REASON_UNSPECIFIED (errorCode is a free-form String on the Kotlin side). @@ -63,6 +64,7 @@ internal object LlmResponseCodec { response.finishReason?.let { builder.finishReason(it.toGenai()) } response.usageMetadata?.let { builder.usageMetadata(UsageMetadataCodec.toJava(it)) } response.groundingMetadata?.let { builder.groundingMetadata(GroundingMetadataCodec.toJava(it)) } + response.cacheMetadata?.let { builder.cacheMetadata(CacheMetadataCodec.toJava(it)) } response.avgLogprobs?.let { builder.avgLogprobs(it) } response.errorCode?.let { builder.errorCode(GenaiFinishReason(it)) } response.customMetadata?.let { builder.customMetadata(CustomMetadataCodec.toJava(it)) } diff --git a/tokt/src/test/kotlin/com/google/adk/tokt/KtRunnerInteropTest.kt b/tokt/src/test/kotlin/com/google/adk/tokt/KtRunnerInteropTest.kt index dd48ff0f0..d47b1f9a8 100644 --- a/tokt/src/test/kotlin/com/google/adk/tokt/KtRunnerInteropTest.kt +++ b/tokt/src/test/kotlin/com/google/adk/tokt/KtRunnerInteropTest.kt @@ -40,6 +40,7 @@ import com.google.adk.kt.apps.App as KtApp import com.google.adk.kt.callbacks.AfterModelCallback as KtAfterModelCallback import com.google.adk.kt.events.Event as KtEvent import com.google.adk.kt.events.EventActions as KtEventActions +import com.google.adk.kt.models.CacheMetadata as KtCacheMetadata import com.google.adk.kt.models.LlmResponse as KtLlmResponse import com.google.adk.kt.runners.InMemoryRunner as KtInMemoryRunner import com.google.adk.kt.runners.Runner as KtRunner @@ -95,6 +96,7 @@ import com.google.adk.kt.workflow.Workflow as KtWorkflow import com.google.adk.memory.InMemoryMemoryService as JavaInMemoryMemoryService import com.google.adk.models.BaseLlm as JavaBaseLlm import com.google.adk.models.BaseLlmConnection as JavaBaseLlmConnection +import com.google.adk.models.CacheMetadata as JavaCacheMetadata import com.google.adk.models.LlmRequest as JavaLlmRequest import com.google.adk.models.LlmResponse as JavaLlmResponse import com.google.adk.plugins.BasePlugin as JavaBasePlugin @@ -3293,6 +3295,52 @@ class KtRunnerInteropTest { ) } + @Test + fun ktRunner_bridgedPluginReturningTheResponse_keepsItsCacheMetadata() = runBlocking { + // The event keeps the response's cache metadata; the next turn finds its cache there. + val model = + object : JavaBaseLlm("java-model") { + override fun generateContent( + llmRequest: JavaLlmRequest, + stream: Boolean, + ): Flowable = + Flowable.just( + JavaLlmResponse.builder() + .content(modelText("cached")) + .cacheMetadata( + JavaCacheMetadata.builder().fingerprint("abc").contentsCount(2).build() + ) + .build() + ) + + override fun connect(llmRequest: JavaLlmRequest): JavaBaseLlmConnection = + throw UnsupportedOperationException() + } + // Returning the response sends it through LlmResponseCodec to Java and back. + val passthrough = + object : JavaBasePlugin("passthrough") { + override fun afterModelCallback( + callbackContext: JavaCallbackContext, + llmResponse: JavaLlmResponse, + ): Maybe = Maybe.just(llmResponse) + } + val runner = + KtInMemoryRunner( + KtApp( + appName = "app", + rootAgent = KtLlmAgent(name = "a", model = JavaAdkToKt.asKtModel(model)), + plugins = listOf(JavaAdkToKt.asKtPlugin(passthrough)), + ) + ) + + val events = runner.turn() + + assertEquals( + KtCacheMetadata(fingerprint = "abc", contentsCount = 2), + events.single { it.author == "a" }.cacheMetadata, + ) + } + @Test fun ktRunner_bridgedNoOpPlugin_preservesAgentTransfer() = runBlocking { // Regression: bridging any Java plugin must not drop the request's Kotlin-only toolsDict in @@ -4310,6 +4358,31 @@ class KtRunnerInteropTest { ) } + @Test + fun asKtSessionService_cacheMetadata_survivesTheJavaStore() = runBlocking { + // The next turn finds the context cache to reuse in the stored events. + val sessions = JavaAdkToKt.asKtSessionService(JavaInMemorySessionService()) + val key = KtSessionKey("app", "u", "s") + val session = sessions.createSession(key) + val fingerprintOnly = KtCacheMetadata(fingerprint = "abc", contentsCount = 2) + val active = + KtCacheMetadata( + fingerprint = "abc", + contentsCount = 2, + cacheName = "cachedContents/1", + expireTime = 2_000_000_000_000L, + invocationsUsed = 3, + createdAt = 1_999_999_000_000L, + ) + + for (cache in listOf(fingerprintOnly, active)) { + val unused = sessions.appendEvent(session, KtEvent(author = "model", cacheMetadata = cache)) + } + + val stored = assertNotNull(sessions.getSession(key)).events.map { it.cacheMetadata } + assertEquals(listOf(fingerprintOnly, active), stored) + } + @Test fun eventCodec_workflowFields_roundTripThroughTheJavaEvent() { // Java code reads these fields too, so a Kotlin Content in the output crosses as a genai one.