feat: Setting up data structures for pause/resume/rewind

PiperOrigin-RevId: 860852195

Author: google-genai-bot

core/src/main/java/com/google/adk/agents/BaseAgentState.java

@@ -0,0 +1,39 @@
+/*
+ * 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.agents;
+
+import com.google.adk.JsonBaseModel;
+
+/** Base class for all agent states. */
+public class BaseAgentState extends JsonBaseModel {
+
+  protected BaseAgentState() {}
+
+  /** Returns a new {@link Builder} for creating {@link BaseAgentState} instances. */
+  public static Builder builder() {
+    return new Builder();
+  }
+
+  /** Builder for {@link BaseAgentState}. */
+  public static class Builder {
+    private Builder() {}
+
+    public BaseAgentState build() {
+      return new BaseAgentState();
+    }
+  }
+}

core/src/main/java/com/google/adk/agents/InvocationContext.java

@@ -50,6 +50,8 @@ public class InvocationContext {
   private final Session session;
   private final Optional<Content> userContent;
   private final RunConfig runConfig;
+  private final Map<String, BaseAgentState> agentStates;
+  private final Map<String, Boolean> endOfAgents;
   private final ResumabilityConfig resumabilityConfig;
   private final InvocationCostManager invocationCostManager;
 
@@ -71,6 +73,8 @@ protected InvocationContext(Builder builder) {
     this.userContent = builder.userContent;
     this.runConfig = builder.runConfig;
     this.endInvocation = builder.endInvocation;
+    this.agentStates = builder.agentStates;
+    this.endOfAgents = builder.endOfAgents;
     this.resumabilityConfig = builder.resumabilityConfig;
     this.invocationCostManager = builder.invocationCostManager;
   }
@@ -299,6 +303,16 @@ public RunConfig runConfig() {
     return runConfig;
   }
 
+  /** Returns agent-specific state saved within this invocation. */
+  public Map<String, BaseAgentState> agentStates() {
+    return agentStates;
+  }
+
+  /** Returns map of agents that ended during this invocation. */
+  public Map<String, Boolean> endOfAgents() {
+    return endOfAgents;
+  }
+
   /**
    * Returns whether this invocation should be ended, e.g., due to reaching a terminal state or
    * error.
@@ -410,6 +424,8 @@ private Builder(InvocationContext context) {
       this.userContent = context.userContent;
       this.runConfig = context.runConfig;
       this.endInvocation = context.endInvocation;
+      this.agentStates = new ConcurrentHashMap<>(context.agentStates);
+      this.endOfAgents = new ConcurrentHashMap<>(context.endOfAgents);
       this.resumabilityConfig = context.resumabilityConfig;
       this.invocationCostManager = context.invocationCostManager;
     }
@@ -427,6 +443,8 @@ private Builder(InvocationContext context) {
     private Optional<Content> userContent = Optional.empty();
     private RunConfig runConfig = RunConfig.builder().build();
     private boolean endInvocation = false;
+    private Map<String, BaseAgentState> agentStates = new ConcurrentHashMap<>();
+    private Map<String, Boolean> endOfAgents = new ConcurrentHashMap<>();
     private ResumabilityConfig resumabilityConfig = new ResumabilityConfig();
     private InvocationCostManager invocationCostManager = new InvocationCostManager();
 
@@ -616,6 +634,30 @@ public Builder endInvocation(boolean endInvocation) {
       return this;
     }
 
+    /**
+     * Sets agent-specific state saved within this invocation.
+     *
+     * @param agentStates agent-specific state saved within this invocation.
+     * @return this builder instance for chaining.
+     */
+    @CanIgnoreReturnValue
+    public Builder agentStates(Map<String, BaseAgentState> agentStates) {
+      this.agentStates = agentStates;
+      return this;
+    }
+
+    /**
+     * Sets agent end-of-invocation status.
+     *
+     * @param endOfAgents agent end-of-invocation status.
+     * @return this builder instance for chaining.
+     */
+    @CanIgnoreReturnValue
+    public Builder endOfAgents(Map<String, Boolean> endOfAgents) {
+      this.endOfAgents = endOfAgents;
+      return this;
+    }
+
     /**
      * Sets the resumability configuration for the current agent run.
      *
@@ -660,6 +702,8 @@ public boolean equals(Object o) {
         && Objects.equals(session, that.session)
         && Objects.equals(userContent, that.userContent)
         && Objects.equals(runConfig, that.runConfig)
+        && Objects.equals(agentStates, that.agentStates)
+        && Objects.equals(endOfAgents, that.endOfAgents)
         && Objects.equals(resumabilityConfig, that.resumabilityConfig)
         && Objects.equals(invocationCostManager, that.invocationCostManager);
   }
@@ -680,6 +724,8 @@ public int hashCode() {
         userContent,
         runConfig,
         endInvocation,
+        agentStates,
+        endOfAgents,
         resumabilityConfig,
         invocationCostManager);
   }

core/src/main/java/com/google/adk/agents/LlmAgent.java

@@ -704,16 +704,7 @@ private static boolean isThought(Part part) {
 
   @Override
   protected Flowable<Event> runAsyncImpl(InvocationContext invocationContext) {
-    return llmFlow
-        .run(invocationContext)
-        .concatMap(
-            event -> {
-              this.maybeSaveOutputToState(event);
-              if (invocationContext.shouldPauseInvocation(event)) {
-                return Flowable.just(event).concatWith(Flowable.empty());
-              }
-              return Flowable.just(event);
-            });
+    return llmFlow.run(invocationContext).doOnNext(this::maybeSaveOutputToState);
   }
 
   @Override

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

@@ -17,6 +17,7 @@
 
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
+import com.google.adk.agents.BaseAgentState;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import com.google.genai.types.Part;
 import java.util.Objects;
@@ -37,8 +38,12 @@ public class EventActions {
   private Optional<Boolean> escalate;
   private ConcurrentMap<String, ConcurrentMap<String, Object>> requestedAuthConfigs;
   private ConcurrentMap<String, ToolConfirmation> requestedToolConfirmations;
+  private boolean endOfAgent;
+  private ConcurrentMap<String, BaseAgentState> agentStates;
+  private boolean pauseInvocation;
   private Optional<Boolean> endInvocation;
   private Optional<EventCompaction> compaction;
+  private Optional<String> rewindBeforeInvocationId;
 
   /** Default constructor for Jackson. */
   public EventActions() {
@@ -49,8 +54,12 @@ public EventActions() {
     this.escalate = Optional.empty();
     this.requestedAuthConfigs = new ConcurrentHashMap<>();
     this.requestedToolConfirmations = new ConcurrentHashMap<>();
+    this.endOfAgent = false;
+    this.pauseInvocation = false;
     this.endInvocation = Optional.empty();
     this.compaction = Optional.empty();
+    this.agentStates = new ConcurrentHashMap<>();
+    this.rewindBeforeInvocationId = Optional.empty();
   }
 
   private EventActions(Builder builder) {
@@ -61,8 +70,12 @@ private EventActions(Builder builder) {
     this.escalate = builder.escalate;
     this.requestedAuthConfigs = builder.requestedAuthConfigs;
     this.requestedToolConfirmations = builder.requestedToolConfirmations;
+    this.endOfAgent = builder.endOfAgent;
+    this.pauseInvocation = builder.pauseInvocation;
     this.endInvocation = builder.endInvocation;
     this.compaction = builder.compaction;
+    this.agentStates = builder.agentStates;
+    this.rewindBeforeInvocationId = builder.rewindBeforeInvocationId;
   }
 
   @JsonProperty("skipSummarization")
@@ -146,6 +159,24 @@ public void setRequestedToolConfirmations(
     this.requestedToolConfirmations = requestedToolConfirmations;
   }
 
+  @JsonProperty("endOfAgent")
+  public boolean endOfAgent() {
+    return endOfAgent;
+  }
+
+  public void setEndOfAgent(boolean endOfAgent) {
+    this.endOfAgent = endOfAgent;
+  }
+
+  @JsonProperty("pauseInvocation")
+  public boolean pauseInvocation() {
+    return pauseInvocation;
+  }
+
+  public void setPauseInvocation(boolean pauseInvocation) {
+    this.pauseInvocation = pauseInvocation;
+  }
+
   @JsonProperty("endInvocation")
   public Optional<Boolean> endInvocation() {
     return endInvocation;
@@ -168,6 +199,24 @@ public void setCompaction(Optional<EventCompaction> compaction) {
     this.compaction = compaction;
   }
 
+  @JsonProperty("agentStates")
+  public ConcurrentMap<String, BaseAgentState> agentStates() {
+    return agentStates;
+  }
+
+  public void setAgentStates(ConcurrentMap<String, BaseAgentState> agentStates) {
+    this.agentStates = agentStates;
+  }
+
+  @JsonProperty("rewindBeforeInvocationId")
+  public Optional<String> rewindBeforeInvocationId() {
+    return rewindBeforeInvocationId;
+  }
+
+  public void setRewindBeforeInvocationId(@Nullable String rewindBeforeInvocationId) {
+    this.rewindBeforeInvocationId = Optional.ofNullable(rewindBeforeInvocationId);
+  }
+
   public static Builder builder() {
     return new Builder();
   }
@@ -191,8 +240,12 @@ public boolean equals(Object o) {
         && Objects.equals(escalate, that.escalate)
         && Objects.equals(requestedAuthConfigs, that.requestedAuthConfigs)
         && Objects.equals(requestedToolConfirmations, that.requestedToolConfirmations)
+        && (endOfAgent == that.endOfAgent)
+        && (pauseInvocation == that.pauseInvocation)
         && Objects.equals(endInvocation, that.endInvocation)
-        && Objects.equals(compaction, that.compaction);
+        && Objects.equals(compaction, that.compaction)
+        && Objects.equals(agentStates, that.agentStates)
+        && Objects.equals(rewindBeforeInvocationId, that.rewindBeforeInvocationId);
   }
 
   @Override
@@ -205,8 +258,12 @@ public int hashCode() {
         escalate,
         requestedAuthConfigs,
         requestedToolConfirmations,
+        endOfAgent,
+        pauseInvocation,
         endInvocation,
-        compaction);
+        compaction,
+        agentStates,
+        rewindBeforeInvocationId);
   }
 
   /** Builder for {@link EventActions}. */
@@ -218,8 +275,12 @@ public static class Builder {
     private Optional<Boolean> escalate;
     private ConcurrentMap<String, ConcurrentMap<String, Object>> requestedAuthConfigs;
     private ConcurrentMap<String, ToolConfirmation> requestedToolConfirmations;
+    private boolean endOfAgent = false;
+    private boolean pauseInvocation = false;
     private Optional<Boolean> endInvocation;
     private Optional<EventCompaction> compaction;
+    private ConcurrentMap<String, BaseAgentState> agentStates;
+    private Optional<String> rewindBeforeInvocationId;
 
     public Builder() {
       this.skipSummarization = Optional.empty();
@@ -231,6 +292,8 @@ public Builder() {
       this.requestedToolConfirmations = new ConcurrentHashMap<>();
       this.endInvocation = Optional.empty();
       this.compaction = Optional.empty();
+      this.agentStates = new ConcurrentHashMap<>();
+      this.rewindBeforeInvocationId = Optional.empty();
     }
 
     private Builder(EventActions eventActions) {
@@ -242,8 +305,12 @@ private Builder(EventActions eventActions) {
       this.requestedAuthConfigs = new ConcurrentHashMap<>(eventActions.requestedAuthConfigs());
       this.requestedToolConfirmations =
           new ConcurrentHashMap<>(eventActions.requestedToolConfirmations());
+      this.endOfAgent = eventActions.endOfAgent();
+      this.pauseInvocation = eventActions.pauseInvocation();
       this.endInvocation = eventActions.endInvocation();
       this.compaction = eventActions.compaction();
+      this.agentStates = new ConcurrentHashMap<>(eventActions.agentStates());
+      this.rewindBeforeInvocationId = eventActions.rewindBeforeInvocationId();
     }
 
     @CanIgnoreReturnValue
@@ -296,6 +363,20 @@ public Builder requestedToolConfirmations(ConcurrentMap<String, ToolConfirmation
       return this;
     }
 
+    @CanIgnoreReturnValue
+    @JsonProperty("endOfAgent")
+    public Builder endOfAgent(boolean endOfAgent) {
+      this.endOfAgent = endOfAgent;
+      return this;
+    }
+
+    @CanIgnoreReturnValue
+    @JsonProperty("pauseInvocation")
+    public Builder pauseInvocation(boolean pauseInvocation) {
+      this.pauseInvocation = pauseInvocation;
+      return this;
+    }
+
     @CanIgnoreReturnValue
     @JsonProperty("endInvocation")
     public Builder endInvocation(boolean endInvocation) {
@@ -310,6 +391,20 @@ public Builder compaction(EventCompaction value) {
       return this;
     }
 
+    @CanIgnoreReturnValue
+    @JsonProperty("agentStates")
+    public Builder agentStates(ConcurrentMap<String, BaseAgentState> agentStates) {
+      this.agentStates = agentStates;
+      return this;
+    }
+
+    @CanIgnoreReturnValue
+    @JsonProperty("rewindBeforeInvocationId")
+    public Builder rewindBeforeInvocationId(String rewindBeforeInvocationId) {
+      this.rewindBeforeInvocationId = Optional.ofNullable(rewindBeforeInvocationId);
+      return this;
+    }
+
     @CanIgnoreReturnValue
     public Builder merge(EventActions other) {
       other.skipSummarization().ifPresent(this::skipSummarization);
@@ -319,8 +414,12 @@ public Builder merge(EventActions other) {
       other.escalate().ifPresent(this::escalate);
       this.requestedAuthConfigs.putAll(other.requestedAuthConfigs());
       this.requestedToolConfirmations.putAll(other.requestedToolConfirmations());
+      this.endOfAgent = other.endOfAgent();
+      this.pauseInvocation = other.pauseInvocation();
       other.endInvocation().ifPresent(this::endInvocation);
       other.compaction().ifPresent(this::compaction);
+      this.agentStates.putAll(other.agentStates());
+      other.rewindBeforeInvocationId().ifPresent(this::rewindBeforeInvocationId);
       return this;
     }