Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
@@ -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(),
)
}
9 changes: 6 additions & 3 deletions tokt/src/main/kotlin/com/google/adk/tokt/codecs/EventCodec.kt
Original file line number Diff line number Diff line change
Expand Up @@ -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 {

Expand All @@ -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)) }
Expand Down Expand Up @@ -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) },
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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 {

Expand All @@ -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).
Expand All @@ -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)) }
Expand Down
73 changes: 73 additions & 0 deletions tokt/src/test/kotlin/com/google/adk/tokt/KtRunnerInteropTest.kt
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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<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
Expand Down Expand Up @@ -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.
Expand Down
Loading