feat: add rewindBeforeInvocationId to EventActions and carry it through the Kotlin interop

PiperOrigin-RevId: 988179449

Author: damianmomotgoogle

core/src/main/java/com/google/adk/events/EventActions.java

@@ -46,6 +46,7 @@ public class EventActions extends JsonBaseModel {
   private @Nullable Map<String, Object> agentState;
   private @Nullable EventCompaction compaction;
   private @Nullable Object setModelResponse;
+  private @Nullable String rewindBeforeInvocationId;
 
   /** Default constructor for Jackson. */
   public EventActions() {
@@ -70,6 +71,7 @@ private EventActions(Builder builder) {
     this.agentState = builder.agentState;
     this.compaction = builder.compaction;
     this.setModelResponse = builder.setModelResponse;
+    this.rewindBeforeInvocationId = builder.rewindBeforeInvocationId;
   }
 
   @JsonProperty("skipSummarization")
@@ -230,6 +232,20 @@ public void setSetModelResponse(@Nullable Object setModelResponse) {
     this.setModelResponse = setModelResponse;
   }
 
+  /**
+   * On a rewind event, the ID of the invocation to rewind before: the history from that invocation
+   * up to the rewind event is dropped. ADK Java does not apply rewinds itself; the field exists for
+   * runtimes that do, such as ADK Kotlin.
+   */
+  @JsonProperty("rewindBeforeInvocationId")
+  public Optional<String> rewindBeforeInvocationId() {
+    return Optional.ofNullable(rewindBeforeInvocationId);
+  }
+
+  public void setRewindBeforeInvocationId(@Nullable String rewindBeforeInvocationId) {
+    this.rewindBeforeInvocationId = rewindBeforeInvocationId;
+  }
+
   public static Builder builder() {
     return new Builder();
   }
@@ -257,7 +273,8 @@ public boolean equals(Object o) {
         && (endOfAgent == that.endOfAgent)
         && Objects.equals(agentState, that.agentState)
         && Objects.equals(compaction, that.compaction)
-        && Objects.equals(setModelResponse, that.setModelResponse);
+        && Objects.equals(setModelResponse, that.setModelResponse)
+        && Objects.equals(rewindBeforeInvocationId, that.rewindBeforeInvocationId);
   }
 
   @Override
@@ -274,7 +291,8 @@ public int hashCode() {
         endOfAgent,
         agentState,
         compaction,
-        setModelResponse);
+        setModelResponse,
+        rewindBeforeInvocationId);
   }
 
   /** Builder for {@link EventActions}. */
@@ -291,6 +309,7 @@ public static class Builder {
     private @Nullable Map<String, Object> agentState;
     private @Nullable EventCompaction compaction;
     private @Nullable Object setModelResponse;
+    private @Nullable String rewindBeforeInvocationId;
 
     public Builder() {
       this.stateDelta = new ConcurrentHashMap<>();
@@ -314,6 +333,7 @@ private Builder(EventActions eventActions) {
       this.agentState = eventActions.agentState;
       this.compaction = eventActions.compaction;
       this.setModelResponse = eventActions.setModelResponse;
+      this.rewindBeforeInvocationId = eventActions.rewindBeforeInvocationId;
     }
 
     @CanIgnoreReturnValue
@@ -434,6 +454,13 @@ public Builder setModelResponse(@Nullable Object value) {
       return this;
     }
 
+    @CanIgnoreReturnValue
+    @JsonProperty("rewindBeforeInvocationId")
+    public Builder rewindBeforeInvocationId(@Nullable String value) {
+      this.rewindBeforeInvocationId = value;
+      return this;
+    }
+
     @CanIgnoreReturnValue
     public Builder merge(EventActions other) {
       other.skipSummarization().ifPresent(this::skipSummarization);
@@ -448,6 +475,7 @@ public Builder merge(EventActions other) {
       other.agentState().ifPresent(this::agentState);
       other.compaction().ifPresent(this::compaction);
       other.setModelResponse().ifPresent(this::setModelResponse);
+      other.rewindBeforeInvocationId().ifPresent(this::rewindBeforeInvocationId);
       return this;
     }
 

core/src/test/java/com/google/adk/events/EventActionsTest.java

@@ -90,6 +90,7 @@ public void merge_mergesAllFields() {
                 new ConcurrentHashMap<>(ImmutableMap.of("tool2", TOOL_CONFIRMATION)))
             .endOfAgent(true)
             .setModelResponse(ImmutableMap.of("field1", "value1"))
+            .rewindBeforeInvocationId("inv1")
             .build();
 
     EventActions merged = eventActions1.toBuilder().merge(eventActions2).build();
@@ -111,6 +112,58 @@ public void merge_mergesAllFields() {
     assertThat(merged.endOfAgent()).isTrue();
     assertThat(merged.compaction()).hasValue(COMPACTION);
     assertThat(merged.setModelResponse()).hasValue(ImmutableMap.of("field1", "value1"));
+    assertThat(merged.rewindBeforeInvocationId()).hasValue("inv1");
+  }
+
+  @Test
+  public void rewindBeforeInvocationId_roundTripsThroughToBuilder() {
+    EventActions actions = EventActions.builder().rewindBeforeInvocationId("inv1").build();
+
+    EventActions rebuilt = actions.toBuilder().build();
+
+    assertThat(rebuilt).isEqualTo(actions);
+    assertThat(rebuilt).isNotEqualTo(EventActions.builder().build());
+    assertThat(rebuilt.rewindBeforeInvocationId()).hasValue("inv1");
+  }
+
+  @Test
+  public void rewindBeforeInvocationId_roundTripsThroughJson() {
+    EventActions actions = EventActions.builder().rewindBeforeInvocationId("inv1").build();
+
+    String json = actions.toJson();
+    EventActions deserialized = EventActions.fromJsonString(json, EventActions.class);
+
+    assertThat(json).contains("\"rewindBeforeInvocationId\":\"inv1\"");
+    assertThat(deserialized.rewindBeforeInvocationId()).hasValue("inv1");
+  }
+
+  @Test
+  public void rewindBeforeInvocationId_absentByDefault_andOmittedFromJson() {
+    EventActions actions = EventActions.builder().build();
+
+    assertThat(actions.rewindBeforeInvocationId()).isEmpty();
+    assertThat(actions.toJson()).doesNotContain("rewindBeforeInvocationId");
+  }
+
+  @Test
+  public void merge_rewindBeforeInvocationId_absentInOther_keepsThis() {
+    EventActions first = EventActions.builder().rewindBeforeInvocationId("inv1").build();
+    EventActions second = EventActions.builder().build();
+
+    EventActions merged = first.toBuilder().merge(second).build();
+
+    assertThat(merged.rewindBeforeInvocationId()).hasValue("inv1");
+  }
+
+  @Test
+  public void setRewindBeforeInvocationId_setsAndClears() {
+    EventActions actions = new EventActions();
+
+    actions.setRewindBeforeInvocationId("inv1");
+    assertThat(actions.rewindBeforeInvocationId()).hasValue("inv1");
+
+    actions.setRewindBeforeInvocationId(null);
+    assertThat(actions.rewindBeforeInvocationId()).isEmpty();
   }
 
   @Test

tokt/src/main/kotlin/com/google/adk/tokt/JavaAdkToKt.kt

@@ -65,8 +65,8 @@ import kotlinx.coroutines.CoroutineDispatcher
  *   `EventActions`, the adapter throws once the Java call returns - the engine's event actions have
  *   no equivalent.
  *
- * Writes to the context's `skipSummarization`, `requestedToolConfirmations`, and `agentState` do
- * cross to the engine.
+ * Writes to the context's `skipSummarization`, `requestedToolConfirmations`, `agentState`, and
+ * `rewindBeforeInvocationId` do cross to the engine.
  */
 object JavaAdkToKt {
 
@@ -139,9 +139,9 @@ object JavaAdkToKt {
   /**
    * Adapts an ADK Java session service for the Kotlin engine, running its calls on `dispatcher`. A
    * Java view of a Kotlin service is unwrapped to that Kotlin service, with no `dispatcher` hop.
-   * Resumption works across the adapter (`EventActions.agentState` crosses), but a rewind reverts
-   * only state and artifacts: ADK Java's `EventActions` has no `rewindBeforeInvocationId`, so the
-   * rewound turns stay in the model's history.
+   * Resumption works across the adapter (`EventActions.agentState` crosses);
+   * `rewindBeforeInvocationId` crosses too, but a rewind drops the rewound turns only if the Java
+   * service stores it, which ADK Java's `VertexAiSessionService` does not.
    */
   @JvmStatic
   @JvmOverloads

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

@@ -26,8 +26,9 @@ import com.google.adk.tokt.codecs.reconcileRemovedSentinels
 /**
  * Copies what a bridged tool or plugin callback wrote on its [JavaEventActions] onto the live
  * [KtEventActions] the engine reads. Fields the live view already writes through are re-applied
- * idempotently; skip-summarization, tool confirmations, agentState and compaction are carried here.
- * A signal the engine cannot represent (requestedAuthConfigs, deletedArtifactIds) throws.
+ * idempotently; skip-summarization, tool confirmations, agentState, compaction and
+ * rewindBeforeInvocationId are carried here. A signal the engine cannot represent
+ * (requestedAuthConfigs, deletedArtifactIds) throws.
  */
 internal fun reconcileActionsToKt(java: JavaEventActions, kt: KtEventActions) {
   require(java.requestedAuthConfigs().isNullOrEmpty()) {
@@ -41,6 +42,7 @@ internal fun reconcileActionsToKt(java: JavaEventActions, kt: KtEventActions) {
   java.skipSummarization().ifPresent { kt.skipSummarization = it }
   java.agentState().ifPresent { kt.agentState = agentStateFromJava(it) }
   java.compaction().ifPresent { kt.compaction = EventCompactionCodec.fromJava(it) }
+  java.rewindBeforeInvocationId().ifPresent { kt.rewindBeforeInvocationId = it }
   kt.endOfAgent = kt.endOfAgent || java.endOfAgent()
   // Merge deltas only when the tool replaced the actions (else they are the same live Kotlin map).
   if (java.stateDelta() !== kt.stateDelta) kt.stateDelta.putAll(java.stateDelta())

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

@@ -31,20 +31,15 @@ import kotlin.jvm.optionals.getOrNull
  *
  * Carries the identity (`id`, `invocationId`, `author`, `branch`), the [content][ContentCodec], the
  * [actions][KtEventActions] (state/artifact deltas, the control-flow flags - transfer, escalate,
- * skipSummarization, endOfAgent - the requested tool confirmations, and the context-compaction
- * summary), the streaming/error signals (`partial`, `turnComplete`, `interrupted`, `errorMessage`,
- * `errorCode`, `finishReason`), the token `usageMetadata` ([UsageMetadataCodec]) and
- * `groundingMetadata` ([GroundingMetadataCodec]), long-running tool ids, and the timestamp.
+ * skipSummarization, endOfAgent - the requested tool confirmations, the context-compaction summary,
+ * the resumable-run `agentState`, and the rewind marker `rewindBeforeInvocationId`), the
+ * streaming/error signals (`partial`, `turnComplete`, `interrupted`, `errorMessage`, `errorCode`,
+ * `finishReason`), the token `usageMetadata` ([UsageMetadataCodec]) and `groundingMetadata`
+ * ([GroundingMetadataCodec]), long-running tool ids, and the timestamp.
  */
 internal object EventCodec {
 
-  /**
-   * Returns the Java [JavaEvent] view of the Kotlin [event].
-   *
-   * `EventActions.agentState` crosses (see [agentStateToJava]), so a resumable run's state survives
-   * a round trip through a Java session service. `rewindBeforeInvocationId` is still dropped: ADK
-   * Java's `EventActions` has no such field.
-   */
+  /** Returns the Java [JavaEvent] view of the Kotlin [event]. */
   fun toJava(event: KtEvent): JavaEvent {
     val builder =
       JavaEvent.builder()
@@ -98,6 +93,7 @@ internal object EventCodec {
           compaction =
             javaActions.compaction().getOrNull()?.let { EventCompactionCodec.fromJava(it) },
           agentState = agentStateFromJava(javaActions.agentState().getOrNull()),
+          rewindBeforeInvocationId = javaActions.rewindBeforeInvocationId().getOrNull(),
         ),
       branch = event.branch().getOrNull(),
       partial = event.partial().getOrNull() ?: false,

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

@@ -73,6 +73,13 @@ internal class KtEventActionsToJavaView(private val actions: KtEventActions) : J
     actions.compaction = compaction?.let { EventCompactionCodec.fromJava(it) }
   }
 
+  override fun rewindBeforeInvocationId(): Optional<String> =
+    Optional.ofNullable(actions.rewindBeforeInvocationId)
+
+  override fun setRewindBeforeInvocationId(rewindBeforeInvocationId: String?) {
+    actions.rewindBeforeInvocationId = rewindBeforeInvocationId
+  }
+
   // Overridden so the write reaches the Kotlin side; the base setters write private fields the
   // overridden getters never read.
 
@@ -104,12 +111,10 @@ internal class KtEventActionsToJavaView(private val actions: KtEventActions) : J
 /**
  * Builds a read-only Java [JavaEventActions] snapshot from a Kotlin [KtEventActions], carrying the
  * deltas, the control-flow signals (skip-summarization, transfer, escalate, end-of-agent), the
- * requested tool confirmations, and the resumable-run `agentState` ([agentStateToJava]). Used when
- * exposing an existing Kotlin event to a Java runner ([EventCodec.toJava]); unlike
- * [KtEventActionsToJavaView] (a live write sink) nothing writes back through it.
- *
- * Does not carry `rewindBeforeInvocationId`: ADK Java's `EventActions` has no such field, so a
- * rewind request does not survive a round trip through a Java session service.
+ * requested tool confirmations, the compaction summary, the resumable-run `agentState`
+ * ([agentStateToJava]), and the rewind marker (`rewindBeforeInvocationId`). Used when exposing an
+ * existing Kotlin event to a Java runner ([EventCodec.toJava]); unlike [KtEventActionsToJavaView]
+ * (a live write sink), nothing writes back through it.
  */
 internal fun eventActionsToJava(actions: KtEventActions): JavaEventActions =
   JavaEventActions.builder()
@@ -124,6 +129,7 @@ internal fun eventActionsToJava(actions: KtEventActions): JavaEventActions =
     .endOfAgent(actions.endOfAgent)
     .compaction(actions.compaction?.let { EventCompactionCodec.toJava(it) })
     .agentState(agentStateToJava(actions.agentState))
+    .rewindBeforeInvocationId(actions.rewindBeforeInvocationId)
     .build()
 
 /** Copies a Kotlin state delta to Java, mapping the Kotlin [KtState.REMOVED] deletion sentinel. */

tokt/src/test/kotlin/com/google/adk/tokt/KtRunnerInteropTest.kt

@@ -921,6 +921,42 @@ class KtRunnerInteropTest {
       )
     }
 
+  @Test
+  fun ktRunner_rewindThroughAnAdaptedJavaSessionService_dropsTheRewoundTurnFromTheNextPrompt() =
+    runBlocking {
+      // Without the marker surviving the Java store, the engine keeps sending the rewound turn.
+      val javaSessions = JavaInMemorySessionService()
+      val model = SequentialJavaModel(listOf(modelText("ok")))
+      val runner =
+        KtInMemoryRunner(
+          app =
+            KtApp(
+              appName = "app",
+              rootAgent = KtLlmAgent(name = "a", model = JavaAdkToKt.asKtModel(model)),
+            ),
+          sessionService = JavaAdkToKt.asKtSessionService(javaSessions),
+        )
+      runner.turn("keep me")
+      val rewound = assertNotNull(runner.turn("forget me").first().invocationId)
+
+      runner.rewindAsync(userId = "u", sessionId = "s", rewindBeforeInvocationId = rewound)
+      runner.turn("after the rewind")
+
+      val stored =
+        javaSessions.getSession("app", "u", "s", Optional.empty()).blockingGet()
+          ?: fail("the adapted Java service should hold the session")
+      assertTrue(
+        stored.events().any { it.actions().rewindBeforeInvocationId().getOrNull() == rewound },
+        "the stored rewind event should keep its marker",
+      )
+      val lastPrompt =
+        model.requests.last().contents().flatMap { content ->
+          content.parts().getOrNull().orEmpty().mapNotNull { it.text().getOrNull() }
+        }
+      assertContains(lastPrompt, "keep me", "the turn before the rewind should stay in the prompt")
+      assertFalse("forget me" in lastPrompt, "the rewound turn must not reach the model")
+    }
+
   @Test
   fun ktRunner_pluginBeforeToolCallback_firesForNativeKotlinTool_andCanDenyIt() = runBlocking {
     // A native Kotlin tool has no Java form, yet a Java plugin's beforeToolCallback must still
@@ -3491,6 +3527,19 @@ class KtRunnerInteropTest {
     )
   }
 
+  @Test
+  fun eventCodec_rewindBeforeInvocationId_roundTripsThroughTheJavaEvent() {
+    // Wiring check: without both directions, a Java session service drops the rewind marker.
+    val ktEvent =
+      KtEvent(author = "user", actions = KtEventActions(rewindBeforeInvocationId = "inv1"))
+
+    assertEquals(
+      "inv1",
+      EventCodec.fromJava(EventCodec.toJava(ktEvent)).actions.rewindBeforeInvocationId,
+      "rewindBeforeInvocationId must survive EventCodec.toJava -> fromJava",
+    )
+  }
+
   @Test
   fun schemaCodec_carriesEveryFacet_roundTrips() {
     // Every facet the Kotlin Schema models must survive Java -> Kotlin -> Java, or structured
@@ -3917,6 +3966,31 @@ class KtRunnerInteropTest {
     assertTrue(kt.endOfAgent, "setEndInvocation should reach the Kotlin actions")
   }
 
+  @Test
+  fun ktEventActionsView_rewindBeforeInvocationId_readsAndWritesThrough() {
+    // A marker set in place by a Java tool must reach the engine's actions.
+    val kt = KtEventActions(rewindBeforeInvocationId = "inv1")
+    val view = KtEventActionsToJavaView(kt)
+
+    assertEquals(
+      "inv1",
+      view.rewindBeforeInvocationId().getOrNull(),
+      "the view should read through",
+    )
+    view.setRewindBeforeInvocationId("inv2")
+    assertEquals("inv2", kt.rewindBeforeInvocationId, "the view should write through")
+  }
+
+  @Test
+  fun reconcileActions_carriesRewindBeforeInvocationId() {
+    // A Java tool that replaces its actions (setActions) must still carry the marker.
+    val kt = KtEventActions()
+
+    reconcileActionsToKt(JavaEventActions.builder().rewindBeforeInvocationId("inv1").build(), kt)
+
+    assertEquals("inv1", kt.rewindBeforeInvocationId)
+  }
+
   @Test
   fun stateValueFromJava_nullMapsToRemovalSentinel() {
     // ADK Java treats a null state value as a removal; the bridge must not NPE on the runAsync