fix: keep context cache metadata in events and model responses

PiperOrigin-RevId: 1007188274

Author: damianmomotgoogle

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(),
+    )
+}

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) },

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)) }

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
@@ -94,6 +95,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
@@ -3292,6 +3294,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<JavaLlmResponse> =
+          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<JavaLlmResponse> = 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
@@ -4309,6 +4357,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.