fix(tokt): merge a Java plugin's replacement onto the original, as ADK Python does

tokt/src/main/kotlin/com/google/adk/tokt/adapters/JavaPluginToKt.kt

@@ -32,6 +32,8 @@ import com.google.adk.tokt.codecs.ContentCodec
 import com.google.adk.tokt.codecs.EventCodec
 import com.google.adk.tokt.codecs.LlmRequestCodec
 import com.google.adk.tokt.codecs.LlmResponseCodec
+import com.google.adk.tokt.codecs.reapplyJavaEvent
+import com.google.adk.tokt.codecs.reapplyJavaResponse
 import com.google.adk.tokt.context.KtInvocationContextToJavaView
 import com.google.adk.tokt.context.javaAgentView
 import com.google.adk.tokt.context.ktCallbackContextToJava
@@ -97,7 +99,7 @@ internal class JavaPluginToKt(
   override suspend fun onEvent(invocationContext: KtInvocationContext, event: KtEvent): KtEvent {
     val javaContext = KtInvocationContextToJavaView(invocationContext, dispatcher)
     val replacement = onDispatcher { plugin.onEventCallback(javaContext, EventCodec.toJava(event)) }
-    return replacement?.let { EventCodec.fromJava(it) } ?: event
+    return if (replacement != null) reapplyJavaEvent(event, replacement) else event
   }
 
   override suspend fun afterRun(invocationContext: KtInvocationContext) {
@@ -157,7 +159,7 @@ internal class JavaPluginToKt(
       plugin.afterModelCallback(javaContext, LlmResponseCodec.toJava(response))
     }
     reconcileActionsToKt(javaContext.eventActions(), context.eventActions)
-    return if (override != null) LlmResponseCodec.fromJava(override) else response
+    return if (override != null) reapplyJavaResponse(response, override) else response
   }
 
   override suspend fun onModelError(

tokt/src/main/kotlin/com/google/adk/tokt/codecs/EventCodec.kt

@@ -37,7 +37,8 @@ import kotlin.jvm.optionals.getOrNull
  * `usageMetadata` ([UsageMetadataCodec]), `groundingMetadata` ([GroundingMetadataCodec]), context
  * `cacheMetadata` ([CacheMetadataCodec]), `avgLogProbs`, `modelVersion`, `customMetadata`, a
  * workflow node's `output` and `nodeInfo` ([WorkflowCodec]), long-running tool ids, the timestamp,
- * and the `isolationScope`. Every other field is dropped.
+ * and the `isolationScope`. Every other field is dropped; [reapplyJavaEvent] keeps them when a Java
+ * plugin returns an event.
  */
 internal object EventCodec {
 
@@ -131,3 +132,31 @@ internal object EventCodec {
     )
   }
 }
+
+/**
+ * Re-applies the event a Java plugin returned onto the original Kotlin [event]: fields [EventCodec]
+ * converts come from the result, and every other field keeps [event]'s value. `partial`,
+ * `turnComplete`, `interrupted` and `usageMetadata` keep [event]'s values when the result leaves
+ * them unset, and `interactionStatus` stays only while the turn is complete. An `output` the plugin
+ * left as tokt sent it keeps its Kotlin value, such as a `JsonElement`.
+ */
+internal fun reapplyJavaEvent(event: KtEvent, javaResult: JavaEvent): KtEvent {
+  val result = EventCodec.fromJava(javaResult)
+  // fromJava reads unset flags as false, so check the Java values.
+  val turnComplete = javaResult.turnComplete().getOrNull() ?: event.turnComplete
+  return result.copy(
+    partial = javaResult.partial().getOrNull() ?: event.partial,
+    turnComplete = turnComplete,
+    interrupted = javaResult.interrupted().getOrNull() ?: event.interrupted,
+    usageMetadata = result.usageMetadata ?: event.usageMetadata,
+    output =
+      if (javaResult.output().getOrNull() == WorkflowCodec.outputToJava(event.output)) event.output
+      else result.output,
+    citationMetadata = event.citationMetadata,
+    turnCompleteReason = event.turnCompleteReason,
+    interactionStatus = event.interactionStatus.takeIf { turnComplete },
+    liveSessionId = event.liveSessionId,
+    liveSessionResumptionUpdate = event.liveSessionResumptionUpdate,
+    voiceActivity = event.voiceActivity,
+  )
+}

tokt/src/main/kotlin/com/google/adk/tokt/codecs/LlmResponseCodec.kt

@@ -30,7 +30,7 @@ import kotlin.jvm.optionals.getOrNull
  * Carries content, error message, streaming flags, `turnComplete`, the live transcriptions, model
  * version, finish reason, usage metadata, grounding and context cache metadata, average log-probs,
  * error code, and custom metadata. Every other field, such as citation metadata and logprobs
- * results, is dropped.
+ * results, is dropped; [reapplyJavaResponse] keeps them when a Java plugin returns a response.
  */
 internal object LlmResponseCodec {
 
@@ -80,3 +80,31 @@ internal object LlmResponseCodec {
     return builder.build()
   }
 }
+
+/**
+ * Re-applies the response a Java plugin returned onto the original Kotlin [response]: fields
+ * [LlmResponseCodec] converts come from the result, and every other field keeps [response]'s value.
+ * `partial`, `turnComplete` and `usageMetadata` keep [response]'s values when the result leaves
+ * them unset. `interactionStatus` stays only while the turn is complete.
+ */
+internal fun reapplyJavaResponse(
+  response: KtLlmResponse,
+  javaResult: JavaLlmResponse,
+): KtLlmResponse {
+  val result = LlmResponseCodec.fromJava(javaResult)
+  val turnComplete = result.turnComplete ?: response.turnComplete
+  return result.copy(
+    // fromJava reads an unset partial as false, so check the Java value.
+    partial = javaResult.partial().getOrNull() ?: response.partial,
+    turnComplete = turnComplete,
+    usageMetadata = result.usageMetadata ?: response.usageMetadata,
+    citationMetadata = response.citationMetadata,
+    logprobsResult = response.logprobsResult,
+    turnCompleteReason = response.turnCompleteReason,
+    interactionStatus = response.interactionStatus.takeIf { turnComplete == true },
+    liveSessionId = response.liveSessionId,
+    liveSessionResumptionUpdate = response.liveSessionResumptionUpdate,
+    goAway = response.goAway,
+    voiceActivity = response.voiceActivity,
+  )
+}