builder = ImmutableList.builder();
- for (AfterAgentCallbackBase callback : afterAgentCallback) {
- if (callback instanceof AfterAgentCallback afterAgentCallbackInstance) {
- builder.add(afterAgentCallbackInstance);
- } else if (callback instanceof AfterAgentCallbackSync afterAgentCallbackSyncInstance) {
- builder.add(
- (callbackContext) ->
- Maybe.fromOptional(afterAgentCallbackSyncInstance.call(callbackContext)));
- } else {
- logger.warn(
- "Invalid afterAgentCallback callback type: {}. Ignoring this callback.",
- callback.getClass().getName());
- }
- }
- return builder.build();
}
+ return callbacks.stream()
+ .flatMap(
+ callback -> {
+ if (asyncClass.isInstance(callback)) {
+ return Stream.of(asyncClass.cast(callback));
+ } else if (syncClass.isInstance(callback)) {
+ return Stream.of(converter.apply(syncClass.cast(callback)));
+ } else {
+ logger.warn(
+ "Invalid {} callback type: {}. Ignoring this callback.",
+ callbackTypeForLogging,
+ callback.getClass().getName());
+ return Stream.empty();
+ }
+ })
+ .collect(ImmutableList.toImmutableList());
}
private CallbackUtil() {}
diff --git a/core/src/main/java/com/google/adk/agents/ContextCacheConfig.java b/core/src/main/java/com/google/adk/agents/ContextCacheConfig.java
new file mode 100644
index 000000000..084700d54
--- /dev/null
+++ b/core/src/main/java/com/google/adk/agents/ContextCacheConfig.java
@@ -0,0 +1,59 @@
+/*
+ * 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 java.time.Duration;
+
+/**
+ * Configuration for context caching across all agents in an app.
+ *
+ * This configuration enables and controls context caching behavior for all LLM agents in an app.
+ * When this config is present on an app, context caching is enabled for all agents. When absent
+ * (null), context caching is disabled.
+ *
+ *
Context caching can significantly reduce costs and improve response times by reusing
+ * previously processed context across multiple requests.
+ *
+ * @param maxInvocations Maximum number of invocations to reuse the same cache before refreshing it.
+ * Defaults to 10.
+ * @param ttl Time-to-live for cache. Defaults to 1800 seconds (30 minutes).
+ * @param minTokens Minimum estimated request tokens required to enable caching. This compares
+ * against the estimated total tokens of the request (system instruction + tools + contents).
+ * Context cache storage may have cost. Set higher to avoid caching small requests where
+ * overhead may exceed benefits. Defaults to 0.
+ */
+public record ContextCacheConfig(int maxInvocations, Duration ttl, int minTokens) {
+
+ public ContextCacheConfig() {
+ this(10, Duration.ofSeconds(1800), 0);
+ }
+
+ /** Returns TTL as string format for cache creation. */
+ public String getTtlString() {
+ return ttl.getSeconds() + "s";
+ }
+
+ @Override
+ public String toString() {
+ return "ContextCacheConfig(maxInvocations="
+ + maxInvocations
+ + ", ttl="
+ + ttl.getSeconds()
+ + "s, minTokens="
+ + minTokens
+ + ")";
+ }
+}
diff --git a/core/src/main/java/com/google/adk/agents/InvocationContext.java b/core/src/main/java/com/google/adk/agents/InvocationContext.java
index ace00db4c..6457a8ca4 100644
--- a/core/src/main/java/com/google/adk/agents/InvocationContext.java
+++ b/core/src/main/java/com/google/adk/agents/InvocationContext.java
@@ -16,20 +16,19 @@
package com.google.adk.agents;
-import com.google.adk.apps.ResumabilityConfig;
+import static com.google.common.base.Strings.isNullOrEmpty;
+
import com.google.adk.artifacts.BaseArtifactService;
-import com.google.adk.events.Event;
import com.google.adk.memory.BaseMemoryService;
import com.google.adk.models.LlmCallsLimitExceededException;
import com.google.adk.plugins.Plugin;
import com.google.adk.plugins.PluginManager;
import com.google.adk.sessions.BaseSessionService;
import com.google.adk.sessions.Session;
-import com.google.common.collect.ImmutableSet;
+import com.google.adk.summarizer.EventsCompactionConfig;
import com.google.errorprone.annotations.CanIgnoreReturnValue;
import com.google.errorprone.annotations.InlineMe;
import com.google.genai.types.Content;
-import com.google.genai.types.FunctionCall;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
@@ -50,10 +49,10 @@ public class InvocationContext {
private final Session session;
private final Optional userContent;
private final RunConfig runConfig;
- private final Map agentStates;
- private final Map endOfAgents;
- private final ResumabilityConfig resumabilityConfig;
+ @Nullable private final EventsCompactionConfig eventsCompactionConfig;
+ @Nullable private final ContextCacheConfig contextCacheConfig;
private final InvocationCostManager invocationCostManager;
+ private final Map callbackContextData;
private Optional branch;
private BaseAgent agent;
@@ -73,10 +72,10 @@ 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.eventsCompactionConfig = builder.eventsCompactionConfig;
+ this.contextCacheConfig = builder.contextCacheConfig;
this.invocationCostManager = builder.invocationCostManager;
+ this.callbackContextData = new ConcurrentHashMap<>(builder.callbackContextData);
}
/**
@@ -257,10 +256,7 @@ public String invocationId() {
/**
* Sets the [branch] ID for the current invocation. A branch represents a fork in the conversation
* history.
- *
- * @deprecated Use {@link #toBuilder()} and {@link Builder#branch(String)} instead.
*/
- @Deprecated(forRemoval = true)
public void branch(@Nullable String branch) {
this.branch = Optional.ofNullable(branch);
}
@@ -303,14 +299,12 @@ public RunConfig runConfig() {
return runConfig;
}
- /** Returns agent-specific state saved within this invocation. */
- public Map agentStates() {
- return agentStates;
- }
-
- /** Returns map of agents that ended during this invocation. */
- public Map endOfAgents() {
- return endOfAgents;
+ /**
+ * Returns a map for storing temporary context data that can be shared between different parts of
+ * the invocation (e.g., before/on/after model callbacks).
+ */
+ public Map callbackContextData() {
+ return callbackContextData;
}
/**
@@ -351,26 +345,14 @@ public void incrementLlmCallsCount() throws LlmCallsLimitExceededException {
this.invocationCostManager.incrementAndEnforceLlmCallsLimit(this.runConfig);
}
- /** Returns whether the current invocation is resumable. */
- public boolean isResumable() {
- return resumabilityConfig.isResumable();
+ /** Returns the events compaction configuration for the current agent run. */
+ public Optional eventsCompactionConfig() {
+ return Optional.ofNullable(eventsCompactionConfig);
}
- /** Returns whether to pause the invocation right after this [event]. */
- public boolean shouldPauseInvocation(Event event) {
- if (!isResumable()) {
- return false;
- }
-
- var longRunningToolIds = event.longRunningToolIds().orElse(ImmutableSet.of());
- if (longRunningToolIds.isEmpty()) {
- return false;
- }
-
- return event.functionCalls().stream()
- .map(FunctionCall::id)
- .flatMap(Optional::stream)
- .anyMatch(functionCallId -> longRunningToolIds.contains(functionCallId));
+ /** Returns the context cache configuration for the current agent run. */
+ public Optional contextCacheConfig() {
+ return Optional.ofNullable(contextCacheConfig);
}
private static class InvocationCostManager {
@@ -424,10 +406,10 @@ 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.eventsCompactionConfig = context.eventsCompactionConfig;
+ this.contextCacheConfig = context.contextCacheConfig;
this.invocationCostManager = context.invocationCostManager;
+ this.callbackContextData = new ConcurrentHashMap<>(context.callbackContextData);
}
private BaseSessionService sessionService;
@@ -443,10 +425,10 @@ private Builder(InvocationContext context) {
private Optional userContent = Optional.empty();
private RunConfig runConfig = RunConfig.builder().build();
private boolean endInvocation = false;
- private Map agentStates = new ConcurrentHashMap<>();
- private Map endOfAgents = new ConcurrentHashMap<>();
- private ResumabilityConfig resumabilityConfig = new ResumabilityConfig();
+ @Nullable private EventsCompactionConfig eventsCompactionConfig;
+ @Nullable private ContextCacheConfig contextCacheConfig;
private InvocationCostManager invocationCostManager = new InvocationCostManager();
+ private Map callbackContextData = new ConcurrentHashMap<>();
/**
* Sets the session service for managing session state.
@@ -635,38 +617,38 @@ public Builder endInvocation(boolean endInvocation) {
}
/**
- * Sets agent-specific state saved within this invocation.
+ * Sets the events compaction configuration for the current agent run.
*
- * @param agentStates agent-specific state saved within this invocation.
+ * @param eventsCompactionConfig the events compaction configuration.
* @return this builder instance for chaining.
*/
@CanIgnoreReturnValue
- public Builder agentStates(Map agentStates) {
- this.agentStates = agentStates;
+ public Builder eventsCompactionConfig(@Nullable EventsCompactionConfig eventsCompactionConfig) {
+ this.eventsCompactionConfig = eventsCompactionConfig;
return this;
}
/**
- * Sets agent end-of-invocation status.
+ * Sets the context cache configuration for the current agent run.
*
- * @param endOfAgents agent end-of-invocation status.
+ * @param contextCacheConfig the context cache configuration.
* @return this builder instance for chaining.
*/
@CanIgnoreReturnValue
- public Builder endOfAgents(Map endOfAgents) {
- this.endOfAgents = endOfAgents;
+ public Builder contextCacheConfig(@Nullable ContextCacheConfig contextCacheConfig) {
+ this.contextCacheConfig = contextCacheConfig;
return this;
}
/**
- * Sets the resumability configuration for the current agent run.
+ * Sets the callback context data for the invocation.
*
- * @param resumabilityConfig the resumability configuration.
+ * @param callbackContextData the callback context data.
* @return this builder instance for chaining.
*/
@CanIgnoreReturnValue
- public Builder resumabilityConfig(ResumabilityConfig resumabilityConfig) {
- this.resumabilityConfig = resumabilityConfig;
+ public Builder callbackContextData(Map callbackContextData) {
+ this.callbackContextData = callbackContextData;
return this;
}
@@ -675,12 +657,33 @@ public Builder resumabilityConfig(ResumabilityConfig resumabilityConfig) {
*
* @throws IllegalStateException if any required parameters are missing.
*/
- // TODO: b/462183912 - Add validation for required parameters.
public InvocationContext build() {
+ validate(this);
return new InvocationContext(this);
}
}
+ /**
+ * Validates the required parameters fields: invocationId, agent, session, and sessionService.
+ *
+ * @param builder the builder to validate.
+ * @throws IllegalStateException if any required parameters are missing.
+ */
+ private static void validate(Builder builder) {
+ if (isNullOrEmpty(builder.invocationId)) {
+ throw new IllegalStateException("Invocation ID must be non-empty.");
+ }
+ if (builder.agent == null) {
+ throw new IllegalStateException("Agent must be set.");
+ }
+ if (builder.session == null) {
+ throw new IllegalStateException("Session must be set.");
+ }
+ if (builder.sessionService == null) {
+ throw new IllegalStateException("Session service must be set.");
+ }
+ }
+
@Override
public boolean equals(Object o) {
if (this == o) {
@@ -702,10 +705,10 @@ 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);
+ && Objects.equals(eventsCompactionConfig, that.eventsCompactionConfig)
+ && Objects.equals(contextCacheConfig, that.contextCacheConfig)
+ && Objects.equals(invocationCostManager, that.invocationCostManager)
+ && Objects.equals(callbackContextData, that.callbackContextData);
}
@Override
@@ -724,9 +727,9 @@ public int hashCode() {
userContent,
runConfig,
endInvocation,
- agentStates,
- endOfAgents,
- resumabilityConfig,
- invocationCostManager);
+ eventsCompactionConfig,
+ contextCacheConfig,
+ invocationCostManager,
+ callbackContextData);
}
}
diff --git a/core/src/main/java/com/google/adk/agents/LlmAgent.java b/core/src/main/java/com/google/adk/agents/LlmAgent.java
index 444985971..1893fb162 100644
--- a/core/src/main/java/com/google/adk/agents/LlmAgent.java
+++ b/core/src/main/java/com/google/adk/agents/LlmAgent.java
@@ -17,6 +17,7 @@
package com.google.adk.agents;
import static com.google.common.collect.ImmutableList.toImmutableList;
+import static java.util.Objects.requireNonNullElse;
import static java.util.stream.Collectors.joining;
import com.fasterxml.jackson.core.JsonProcessingException;
@@ -61,6 +62,7 @@
import com.google.genai.types.GenerateContentConfig;
import com.google.genai.types.Part;
import com.google.genai.types.Schema;
+import io.reactivex.rxjava3.core.Completable;
import io.reactivex.rxjava3.core.Flowable;
import io.reactivex.rxjava3.core.Maybe;
import io.reactivex.rxjava3.core.Single;
@@ -103,12 +105,12 @@ public enum IncludeContents {
private final Optional maxSteps;
private final boolean disallowTransferToParent;
private final boolean disallowTransferToPeers;
- private final Optional> beforeModelCallback;
- private final Optional> afterModelCallback;
- private final Optional> onModelErrorCallback;
- private final Optional> beforeToolCallback;
- private final Optional> afterToolCallback;
- private final Optional> onToolErrorCallback;
+ private final ImmutableList beforeModelCallback;
+ private final ImmutableList afterModelCallback;
+ private final ImmutableList onModelErrorCallback;
+ private final ImmutableList beforeToolCallback;
+ private final ImmutableList afterToolCallback;
+ private final ImmutableList onToolErrorCallback;
private final Optional inputSchema;
private final Optional outputSchema;
private final Optional executor;
@@ -126,29 +128,28 @@ protected LlmAgent(Builder builder) {
builder.beforeAgentCallback,
builder.afterAgentCallback);
this.model = Optional.ofNullable(builder.model);
- this.instruction =
- builder.instruction == null ? new Instruction.Static("") : builder.instruction;
+ this.instruction = requireNonNullElse(builder.instruction, new Instruction.Static(""));
this.globalInstruction =
- builder.globalInstruction == null ? new Instruction.Static("") : builder.globalInstruction;
+ requireNonNullElse(builder.globalInstruction, new Instruction.Static(""));
this.generateContentConfig = Optional.ofNullable(builder.generateContentConfig);
this.exampleProvider = Optional.ofNullable(builder.exampleProvider);
- this.includeContents =
- builder.includeContents != null ? builder.includeContents : IncludeContents.DEFAULT;
+ this.includeContents = requireNonNullElse(builder.includeContents, IncludeContents.DEFAULT);
this.planning = builder.planning != null && builder.planning;
this.maxSteps = Optional.ofNullable(builder.maxSteps);
this.disallowTransferToParent = builder.disallowTransferToParent;
this.disallowTransferToPeers = builder.disallowTransferToPeers;
- this.beforeModelCallback = Optional.ofNullable(builder.beforeModelCallback);
- this.afterModelCallback = Optional.ofNullable(builder.afterModelCallback);
- this.onModelErrorCallback = Optional.ofNullable(builder.onModelErrorCallback);
- this.beforeToolCallback = Optional.ofNullable(builder.beforeToolCallback);
- this.afterToolCallback = Optional.ofNullable(builder.afterToolCallback);
- this.onToolErrorCallback = Optional.ofNullable(builder.onToolErrorCallback);
+ this.beforeModelCallback = requireNonNullElse(builder.beforeModelCallback, ImmutableList.of());
+ this.afterModelCallback = requireNonNullElse(builder.afterModelCallback, ImmutableList.of());
+ this.onModelErrorCallback =
+ requireNonNullElse(builder.onModelErrorCallback, ImmutableList.of());
+ this.beforeToolCallback = requireNonNullElse(builder.beforeToolCallback, ImmutableList.of());
+ this.afterToolCallback = requireNonNullElse(builder.afterToolCallback, ImmutableList.of());
+ this.onToolErrorCallback = requireNonNullElse(builder.onToolErrorCallback, ImmutableList.of());
this.inputSchema = Optional.ofNullable(builder.inputSchema);
this.outputSchema = Optional.ofNullable(builder.outputSchema);
this.executor = Optional.ofNullable(builder.executor);
this.outputKey = Optional.ofNullable(builder.outputKey);
- this.toolsUnion = builder.toolsUnion != null ? builder.toolsUnion : ImmutableList.of();
+ this.toolsUnion = requireNonNullElse(builder.toolsUnion, ImmutableList.of());
this.toolsets = extractToolsets(this.toolsUnion);
this.codeExecutor = Optional.ofNullable(builder.codeExecutor);
@@ -841,27 +842,27 @@ public boolean disallowTransferToPeers() {
return disallowTransferToPeers;
}
- public Optional> beforeModelCallback() {
+ public List beforeModelCallback() {
return beforeModelCallback;
}
- public Optional> afterModelCallback() {
+ public List afterModelCallback() {
return afterModelCallback;
}
- public Optional> beforeToolCallback() {
+ public List beforeToolCallback() {
return beforeToolCallback;
}
- public Optional> afterToolCallback() {
+ public List afterToolCallback() {
return afterToolCallback;
}
- public Optional> onModelErrorCallback() {
+ public List onModelErrorCallback() {
return onModelErrorCallback;
}
- public Optional> onToolErrorCallback() {
+ public List onToolErrorCallback() {
return onToolErrorCallback;
}
@@ -871,7 +872,7 @@ public Optional> onToolErrorCallback() {
* This method is only for use by Agent Development Kit.
*/
public List canonicalBeforeModelCallbacks() {
- return beforeModelCallback.orElse(ImmutableList.of());
+ return beforeModelCallback;
}
/**
@@ -880,7 +881,7 @@ public List canonicalBeforeModelCallbacks() {
*
This method is only for use by Agent Development Kit.
*/
public List canonicalAfterModelCallbacks() {
- return afterModelCallback.orElse(ImmutableList.of());
+ return afterModelCallback;
}
/**
@@ -889,7 +890,7 @@ public List canonicalAfterModelCallbacks() {
*
This method is only for use by Agent Development Kit.
*/
public List canonicalOnModelErrorCallbacks() {
- return onModelErrorCallback.orElse(ImmutableList.of());
+ return onModelErrorCallback;
}
/**
@@ -898,7 +899,7 @@ public List canonicalOnModelErrorCallbacks() {
*
This method is only for use by Agent Development Kit.
*/
public List canonicalBeforeToolCallbacks() {
- return beforeToolCallback.orElse(ImmutableList.of());
+ return beforeToolCallback;
}
/**
@@ -907,7 +908,7 @@ public List canonicalBeforeToolCallbacks() {
*
This method is only for use by Agent Development Kit.
*/
public List canonicalAfterToolCallbacks() {
- return afterToolCallback.orElse(ImmutableList.of());
+ return afterToolCallback;
}
/**
@@ -916,7 +917,7 @@ public List canonicalAfterToolCallbacks() {
*
This method is only for use by Agent Development Kit.
*/
public List canonicalOnToolErrorCallbacks() {
- return onToolErrorCallback.orElse(ImmutableList.of());
+ return onToolErrorCallback;
}
public Optional inputSchema() {
@@ -935,9 +936,8 @@ public Optional outputKey() {
return outputKey;
}
- @Nullable
- public BaseCodeExecutor codeExecutor() {
- return codeExecutor.orElse(null);
+ public Optional codeExecutor() {
+ return codeExecutor;
}
public Model resolvedModel() {
@@ -1056,6 +1056,26 @@ public static LlmAgent fromConfig(LlmAgentConfig config, String configAbsPath)
return agent;
}
+ @Override
+ public Completable close() {
+ List completables = new ArrayList<>();
+ toolsets()
+ .forEach(
+ toolset ->
+ completables.add(
+ Completable.fromAction(
+ () -> {
+ try {
+ toolset.close();
+ } catch (Exception e) {
+ logger.error("Failed to close toolset", e);
+ throw e;
+ }
+ })));
+ completables.add(super.close());
+ return Completable.mergeDelayError(completables);
+ }
+
private static void setCallbacksFromConfig(LlmAgentConfig config, Builder builder)
throws ConfigurationException {
ConfigAgentUtils.resolveAndSetCallback(
diff --git a/core/src/main/java/com/google/adk/agents/RunConfig.java b/core/src/main/java/com/google/adk/agents/RunConfig.java
index 308169e36..1ca203eaf 100644
--- a/core/src/main/java/com/google/adk/agents/RunConfig.java
+++ b/core/src/main/java/com/google/adk/agents/RunConfig.java
@@ -134,6 +134,9 @@ public abstract Builder setInputAudioTranscription(
public RunConfig build() {
RunConfig runConfig = autoBuild();
+ if (runConfig.maxLlmCalls() == Integer.MAX_VALUE) {
+ throw new IllegalArgumentException("maxLlmCalls should be less than Integer.MAX_VALUE.");
+ }
if (runConfig.maxLlmCalls() < 0) {
logger.warn(
"maxLlmCalls is negative. This will result in no enforcement on total"
diff --git a/core/src/main/java/com/google/adk/apps/App.java b/core/src/main/java/com/google/adk/apps/App.java
index 3b1f0613a..18e8753c7 100644
--- a/core/src/main/java/com/google/adk/apps/App.java
+++ b/core/src/main/java/com/google/adk/apps/App.java
@@ -17,7 +17,8 @@
package com.google.adk.apps;
import com.google.adk.agents.BaseAgent;
-import com.google.adk.plugins.BasePlugin;
+import com.google.adk.agents.ContextCacheConfig;
+import com.google.adk.plugins.Plugin;
import com.google.adk.summarizer.EventsCompactionConfig;
import com.google.common.collect.ImmutableList;
import com.google.errorprone.annotations.CanIgnoreReturnValue;
@@ -38,21 +39,21 @@ public class App {
private final String name;
private final BaseAgent rootAgent;
- private final ImmutableList plugins;
+ private final ImmutableList plugins;
@Nullable private final EventsCompactionConfig eventsCompactionConfig;
- @Nullable private final ResumabilityConfig resumabilityConfig;
+ @Nullable private final ContextCacheConfig contextCacheConfig;
private App(
String name,
BaseAgent rootAgent,
- List plugins,
+ List plugins,
@Nullable EventsCompactionConfig eventsCompactionConfig,
- @Nullable ResumabilityConfig resumabilityConfig) {
+ @Nullable ContextCacheConfig contextCacheConfig) {
this.name = name;
this.rootAgent = rootAgent;
this.plugins = ImmutableList.copyOf(plugins);
this.eventsCompactionConfig = eventsCompactionConfig;
- this.resumabilityConfig = resumabilityConfig;
+ this.contextCacheConfig = contextCacheConfig;
}
public String name() {
@@ -63,7 +64,7 @@ public BaseAgent rootAgent() {
return rootAgent;
}
- public ImmutableList plugins() {
+ public ImmutableList plugins() {
return plugins;
}
@@ -73,17 +74,17 @@ public EventsCompactionConfig eventsCompactionConfig() {
}
@Nullable
- public ResumabilityConfig resumabilityConfig() {
- return resumabilityConfig;
+ public ContextCacheConfig contextCacheConfig() {
+ return contextCacheConfig;
}
/** Builder for {@link App}. */
public static class Builder {
private String name;
private BaseAgent rootAgent;
- private List plugins = ImmutableList.of();
+ private List plugins = ImmutableList.of();
@Nullable private EventsCompactionConfig eventsCompactionConfig;
- @Nullable private ResumabilityConfig resumabilityConfig;
+ @Nullable private ContextCacheConfig contextCacheConfig;
@CanIgnoreReturnValue
public Builder name(String name) {
@@ -98,7 +99,7 @@ public Builder rootAgent(BaseAgent rootAgent) {
}
@CanIgnoreReturnValue
- public Builder plugins(List plugins) {
+ public Builder plugins(List plugins) {
this.plugins = plugins;
return this;
}
@@ -110,8 +111,8 @@ public Builder eventsCompactionConfig(EventsCompactionConfig eventsCompactionCon
}
@CanIgnoreReturnValue
- public Builder resumabilityConfig(ResumabilityConfig resumabilityConfig) {
- this.resumabilityConfig = resumabilityConfig;
+ public Builder contextCacheConfig(ContextCacheConfig contextCacheConfig) {
+ this.contextCacheConfig = contextCacheConfig;
return this;
}
@@ -123,7 +124,7 @@ public App build() {
throw new IllegalStateException("Root agent must be provided.");
}
validateAppName(name);
- return new App(name, rootAgent, plugins, eventsCompactionConfig, resumabilityConfig);
+ return new App(name, rootAgent, plugins, eventsCompactionConfig, contextCacheConfig);
}
}
diff --git a/core/src/main/java/com/google/adk/apps/ResumabilityConfig.java b/core/src/main/java/com/google/adk/apps/ResumabilityConfig.java
deleted file mode 100644
index b80ce709c..000000000
--- a/core/src/main/java/com/google/adk/apps/ResumabilityConfig.java
+++ /dev/null
@@ -1,28 +0,0 @@
-/*
- * Copyright 2025 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 language governing permissions and
- * limitations under the License.
- */
-package com.google.adk.apps;
-
-/**
- * An app contains Resumability configuration for the agents.
- *
- * @param isResumable Whether the app is resumable.
- */
-public record ResumabilityConfig(boolean isResumable) {
-
- /** Creates a new {@code ResumabilityConfig} with resumability disabled. */
- public ResumabilityConfig() {
- this(false);
- }
-}
diff --git a/core/src/main/java/com/google/adk/artifacts/BaseArtifactService.java b/core/src/main/java/com/google/adk/artifacts/BaseArtifactService.java
index 847e88dd9..b6a3cee23 100644
--- a/core/src/main/java/com/google/adk/artifacts/BaseArtifactService.java
+++ b/core/src/main/java/com/google/adk/artifacts/BaseArtifactService.java
@@ -39,6 +39,28 @@ public interface BaseArtifactService {
Single saveArtifact(
String appName, String userId, String sessionId, String filename, Part artifact);
+ /**
+ * Saves an artifact and returns it with fileData if available.
+ *
+ * Implementations should override this default method for efficiency, as the default performs
+ * two I/O operations (save then load).
+ *
+ * @param appName the app name
+ * @param userId the user ID
+ * @param sessionId the session ID
+ * @param filename the filename
+ * @param artifact the artifact to save
+ * @return the saved artifact with fileData if available.
+ */
+ default Single saveAndReloadArtifact(
+ String appName, String userId, String sessionId, String filename, Part artifact) {
+ return saveArtifact(appName, userId, sessionId, filename, artifact)
+ .flatMap(
+ version ->
+ loadArtifact(appName, userId, sessionId, filename, Optional.of(version))
+ .toSingle());
+ }
+
/**
* Gets an artifact.
*
diff --git a/core/src/main/java/com/google/adk/artifacts/GcsArtifactService.java b/core/src/main/java/com/google/adk/artifacts/GcsArtifactService.java
index 1bfef8cf8..b9bc49a02 100644
--- a/core/src/main/java/com/google/adk/artifacts/GcsArtifactService.java
+++ b/core/src/main/java/com/google/adk/artifacts/GcsArtifactService.java
@@ -18,6 +18,7 @@
import static java.util.Collections.max;
+import com.google.auto.value.AutoValue;
import com.google.cloud.storage.Blob;
import com.google.cloud.storage.BlobId;
import com.google.cloud.storage.BlobInfo;
@@ -27,6 +28,7 @@
import com.google.common.base.Splitter;
import com.google.common.base.VerifyException;
import com.google.common.collect.ImmutableList;
+import com.google.genai.types.FileData;
import com.google.genai.types.Part;
import io.reactivex.rxjava3.core.Completable;
import io.reactivex.rxjava3.core.Maybe;
@@ -108,34 +110,8 @@ private String getBlobName(
@Override
public Single saveArtifact(
String appName, String userId, String sessionId, String filename, Part artifact) {
- return listVersions(appName, userId, sessionId, filename)
- .map(versions -> versions.isEmpty() ? 0 : max(versions) + 1)
- .map(
- nextVersion -> {
- String blobName = getBlobName(appName, userId, sessionId, filename, nextVersion);
- BlobId blobId = BlobId.of(bucketName, blobName);
-
- BlobInfo blobInfo =
- BlobInfo.newBuilder(blobId)
- .setContentType(artifact.inlineData().get().mimeType().orElse(null))
- .build();
-
- try {
- byte[] dataToSave =
- artifact
- .inlineData()
- .get()
- .data()
- .orElseThrow(
- () ->
- new IllegalArgumentException(
- "Saveable artifact data must be non-empty."));
- storageClient.create(blobInfo, dataToSave);
- return nextVersion;
- } catch (StorageException e) {
- throw new VerifyException("Failed to save artifact to GCS", e);
- }
- });
+ return saveArtifactAndReturnBlob(appName, userId, sessionId, filename, artifact)
+ .map(SaveResult::version);
}
/**
@@ -275,4 +251,75 @@ public Single> listVersions(
return Single.just(ImmutableList.of());
}
}
+
+ @Override
+ public Single saveAndReloadArtifact(
+ String appName, String userId, String sessionId, String filename, Part artifact) {
+ return saveArtifactAndReturnBlob(appName, userId, sessionId, filename, artifact)
+ .flatMap(
+ blob -> {
+ Blob savedBlob = blob.blob();
+ String resultMimeType =
+ Optional.ofNullable(savedBlob.getContentType())
+ .or(
+ () ->
+ artifact.inlineData().flatMap(com.google.genai.types.Blob::mimeType))
+ .orElse("application/octet-stream");
+ return Single.just(
+ Part.builder()
+ .fileData(
+ FileData.builder()
+ .fileUri("gs://" + savedBlob.getBucket() + "/" + savedBlob.getName())
+ .mimeType(resultMimeType)
+ .build())
+ .build());
+ });
+ }
+
+ @AutoValue
+ abstract static class SaveResult {
+ static SaveResult create(Blob blob, int version) {
+ return new AutoValue_GcsArtifactService_SaveResult(blob, version);
+ }
+
+ abstract Blob blob();
+
+ abstract int version();
+ }
+
+ private Single saveArtifactAndReturnBlob(
+ String appName, String userId, String sessionId, String filename, Part artifact) {
+ return listVersions(appName, userId, sessionId, filename)
+ .map(versions -> versions.isEmpty() ? 0 : max(versions) + 1)
+ .map(
+ nextVersion -> {
+ if (artifact.inlineData().isEmpty()) {
+ throw new IllegalArgumentException("Saveable artifact must have inline data.");
+ }
+
+ String blobName = getBlobName(appName, userId, sessionId, filename, nextVersion);
+ BlobId blobId = BlobId.of(bucketName, blobName);
+
+ BlobInfo blobInfo =
+ BlobInfo.newBuilder(blobId)
+ .setContentType(artifact.inlineData().get().mimeType().orElse(null))
+ .build();
+
+ try {
+ byte[] dataToSave =
+ artifact
+ .inlineData()
+ .get()
+ .data()
+ .orElseThrow(
+ () ->
+ new IllegalArgumentException(
+ "Saveable artifact data must be non-empty."));
+ Blob blob = storageClient.create(blobInfo, dataToSave);
+ return SaveResult.create(blob, nextVersion);
+ } catch (StorageException e) {
+ throw new VerifyException("Failed to save artifact to GCS", e);
+ }
+ });
+ }
}
diff --git a/core/src/main/java/com/google/adk/artifacts/InMemoryArtifactService.java b/core/src/main/java/com/google/adk/artifacts/InMemoryArtifactService.java
index 890820196..5808f7083 100644
--- a/core/src/main/java/com/google/adk/artifacts/InMemoryArtifactService.java
+++ b/core/src/main/java/com/google/adk/artifacts/InMemoryArtifactService.java
@@ -48,11 +48,8 @@ public InMemoryArtifactService() {
public Single saveArtifact(
String appName, String userId, String sessionId, String filename, Part artifact) {
List versions =
- artifacts
- .computeIfAbsent(appName, k -> new HashMap<>())
- .computeIfAbsent(userId, k -> new HashMap<>())
- .computeIfAbsent(sessionId, k -> new HashMap<>())
- .computeIfAbsent(filename, k -> new ArrayList<>());
+ getArtifactsMap(appName, userId, sessionId)
+ .computeIfAbsent(filename, unused -> new ArrayList<>());
versions.add(artifact);
return Single.just(versions.size() - 1);
}
@@ -66,11 +63,8 @@ public Single saveArtifact(
public Maybe loadArtifact(
String appName, String userId, String sessionId, String filename, Optional version) {
List versions =
- artifacts
- .getOrDefault(appName, new HashMap<>())
- .getOrDefault(userId, new HashMap<>())
- .getOrDefault(sessionId, new HashMap<>())
- .getOrDefault(filename, new ArrayList<>());
+ getArtifactsMap(appName, userId, sessionId)
+ .computeIfAbsent(filename, unused -> new ArrayList<>());
if (versions.isEmpty()) {
return Maybe.empty();
@@ -97,13 +91,7 @@ public Single listArtifactKeys(
String appName, String userId, String sessionId) {
return Single.just(
ListArtifactsResponse.builder()
- .filenames(
- ImmutableList.copyOf(
- artifacts
- .getOrDefault(appName, new HashMap<>())
- .getOrDefault(userId, new HashMap<>())
- .getOrDefault(sessionId, new HashMap<>())
- .keySet()))
+ .filenames(ImmutableList.copyOf(getArtifactsMap(appName, userId, sessionId).keySet()))
.build());
}
@@ -115,11 +103,7 @@ public Single listArtifactKeys(
@Override
public Completable deleteArtifact(
String appName, String userId, String sessionId, String filename) {
- artifacts
- .getOrDefault(appName, new HashMap<>())
- .getOrDefault(userId, new HashMap<>())
- .getOrDefault(sessionId, new HashMap<>())
- .remove(filename);
+ getArtifactsMap(appName, userId, sessionId).remove(filename);
return Completable.complete();
}
@@ -132,15 +116,29 @@ public Completable deleteArtifact(
public Single> listVersions(
String appName, String userId, String sessionId, String filename) {
int size =
- artifacts
- .getOrDefault(appName, new HashMap<>())
- .getOrDefault(userId, new HashMap<>())
- .getOrDefault(sessionId, new HashMap<>())
- .getOrDefault(filename, new ArrayList<>())
+ getArtifactsMap(appName, userId, sessionId)
+ .computeIfAbsent(filename, unused -> new ArrayList<>())
.size();
if (size == 0) {
return Single.just(ImmutableList.of());
}
return Single.just(IntStream.range(0, size).boxed().collect(toImmutableList()));
}
+
+ @Override
+ public Single saveAndReloadArtifact(
+ String appName, String userId, String sessionId, String filename, Part artifact) {
+ return saveArtifact(appName, userId, sessionId, filename, artifact)
+ .flatMap(
+ version ->
+ loadArtifact(appName, userId, sessionId, filename, Optional.of(version))
+ .toSingle());
+ }
+
+ private Map> getArtifactsMap(String appName, String userId, String sessionId) {
+ return artifacts
+ .computeIfAbsent(appName, unused -> new HashMap<>())
+ .computeIfAbsent(userId, unused -> new HashMap<>())
+ .computeIfAbsent(sessionId, unused -> new HashMap<>());
+ }
}
diff --git a/core/src/main/java/com/google/adk/codeexecutors/CodeExecutorContext.java b/core/src/main/java/com/google/adk/codeexecutors/CodeExecutorContext.java
index bad83ebc0..a34102225 100644
--- a/core/src/main/java/com/google/adk/codeexecutors/CodeExecutorContext.java
+++ b/core/src/main/java/com/google/adk/codeexecutors/CodeExecutorContext.java
@@ -89,7 +89,8 @@ public void setExecutionId(String sessionId) {
* @return A list of processed file names in the code executor context.
*/
public List getProcessedFileNames() {
- return (List) this.context.getOrDefault(PROCESSED_FILE_NAMES_KEY, new ArrayList<>());
+ return (List)
+ this.context.computeIfAbsent(PROCESSED_FILE_NAMES_KEY, unused -> new ArrayList<>());
}
/**
@@ -100,7 +101,7 @@ public List getProcessedFileNames() {
public void addProcessedFileNames(List fileNames) {
List processedFileNames =
(List)
- this.context.computeIfAbsent(PROCESSED_FILE_NAMES_KEY, k -> new ArrayList<>());
+ this.context.computeIfAbsent(PROCESSED_FILE_NAMES_KEY, unused -> new ArrayList<>());
processedFileNames.addAll(fileNames);
}
@@ -126,7 +127,7 @@ public List getInputFiles() {
public void addInputFiles(List inputFiles) {
List> fileMaps =
(List>)
- this.sessionState.computeIfAbsent(INPUT_FILE_KEY, k -> new ArrayList<>());
+ this.sessionState.computeIfAbsent(INPUT_FILE_KEY, unused -> new ArrayList<>());
for (File inputFile : inputFiles) {
fileMaps.add(
objectMapper.convertValue(inputFile, new TypeReference>() {}));
@@ -166,7 +167,7 @@ public int getErrorCount(String invocationId) {
public void incrementErrorCount(String invocationId) {
Map errorCounts =
(Map)
- this.sessionState.computeIfAbsent(ERROR_COUNT_KEY, k -> new HashMap<>());
+ this.sessionState.computeIfAbsent(ERROR_COUNT_KEY, unused -> new HashMap<>());
errorCounts.put(invocationId, getErrorCount(invocationId) + 1);
}
@@ -176,9 +177,6 @@ public void incrementErrorCount(String invocationId) {
* @param invocationId The invocation ID to reset the error count for.
*/
public void resetErrorCount(String invocationId) {
- if (!this.sessionState.containsKey(ERROR_COUNT_KEY)) {
- return;
- }
Map errorCounts =
(Map) this.sessionState.get(ERROR_COUNT_KEY);
if (errorCounts != null) {
@@ -198,9 +196,10 @@ public void updateCodeExecutionResult(
String invocationId, String code, String resultStdout, String resultStderr) {
Map>> codeExecutionResults =
(Map>>)
- this.sessionState.computeIfAbsent(CODE_EXECUTION_RESULTS_KEY, k -> new HashMap<>());
+ this.sessionState.computeIfAbsent(
+ CODE_EXECUTION_RESULTS_KEY, unused -> new HashMap<>());
List> resultsForInvocation =
- codeExecutionResults.computeIfAbsent(invocationId, k -> new ArrayList<>());
+ codeExecutionResults.computeIfAbsent(invocationId, unused -> new ArrayList<>());
Map newResult = new HashMap<>();
newResult.put("code", code);
newResult.put("result_stdout", resultStdout);
@@ -210,6 +209,7 @@ public void updateCodeExecutionResult(
}
private Map getCodeExecutorContext(Map sessionState) {
- return (Map) sessionState.computeIfAbsent(CONTEXT_KEY, k -> new HashMap<>());
+ return (Map)
+ sessionState.computeIfAbsent(CONTEXT_KEY, unused -> new HashMap<>());
}
}
diff --git a/core/src/main/java/com/google/adk/events/Event.java b/core/src/main/java/com/google/adk/events/Event.java
index 9e05918be..d968efa53 100644
--- a/core/src/main/java/com/google/adk/events/Event.java
+++ b/core/src/main/java/com/google/adk/events/Event.java
@@ -294,8 +294,7 @@ public final boolean hasTrailingCodeExecutionResult() {
/** Returns true if this is a final response. */
@JsonIgnore
public final boolean finalResponse() {
- if (actions().skipSummarization().orElse(false)
- || (longRunningToolIds().isPresent() && !longRunningToolIds().get().isEmpty())) {
+ if (actions().skipSummarization().orElse(false)) {
return true;
}
return functionCalls().isEmpty()
diff --git a/core/src/main/java/com/google/adk/events/EventActions.java b/core/src/main/java/com/google/adk/events/EventActions.java
index 63909ee1a..6d8c698dd 100644
--- a/core/src/main/java/com/google/adk/events/EventActions.java
+++ b/core/src/main/java/com/google/adk/events/EventActions.java
@@ -18,11 +18,13 @@
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
-import com.google.adk.agents.BaseAgentState;
+import com.google.adk.JsonBaseModel;
+import com.google.adk.sessions.State;
import com.google.errorprone.annotations.CanIgnoreReturnValue;
-import com.google.genai.types.Part;
+import java.util.HashSet;
import java.util.Objects;
import java.util.Optional;
+import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import javax.annotation.Nullable;
@@ -30,50 +32,44 @@
/** Represents the actions attached to an event. */
// TODO - b/414081262 make json wire camelCase
@JsonDeserialize(builder = EventActions.Builder.class)
-public class EventActions {
+public class EventActions extends JsonBaseModel {
private Optional skipSummarization;
private ConcurrentMap stateDelta;
- private ConcurrentMap artifactDelta;
+ private ConcurrentMap artifactDelta;
+ private Set deletedArtifactIds;
private Optional transferToAgent;
private Optional escalate;
private ConcurrentMap> requestedAuthConfigs;
private ConcurrentMap requestedToolConfirmations;
private boolean endOfAgent;
- private ConcurrentMap agentState;
- private Optional endInvocation;
private Optional compaction;
- private Optional rewindBeforeInvocationId;
/** Default constructor for Jackson. */
public EventActions() {
this.skipSummarization = Optional.empty();
this.stateDelta = new ConcurrentHashMap<>();
this.artifactDelta = new ConcurrentHashMap<>();
+ this.deletedArtifactIds = new HashSet<>();
this.transferToAgent = Optional.empty();
this.escalate = Optional.empty();
this.requestedAuthConfigs = new ConcurrentHashMap<>();
this.requestedToolConfirmations = new ConcurrentHashMap<>();
this.endOfAgent = false;
- this.endInvocation = Optional.empty();
this.compaction = Optional.empty();
- this.agentState = new ConcurrentHashMap<>();
- this.rewindBeforeInvocationId = Optional.empty();
}
private EventActions(Builder builder) {
this.skipSummarization = builder.skipSummarization;
this.stateDelta = builder.stateDelta;
this.artifactDelta = builder.artifactDelta;
+ this.deletedArtifactIds = builder.deletedArtifactIds;
this.transferToAgent = builder.transferToAgent;
this.escalate = builder.escalate;
this.requestedAuthConfigs = builder.requestedAuthConfigs;
this.requestedToolConfirmations = builder.requestedToolConfirmations;
this.endOfAgent = builder.endOfAgent;
- this.endInvocation = builder.endInvocation;
this.compaction = builder.compaction;
- this.agentState = builder.agentState;
- this.rewindBeforeInvocationId = builder.rewindBeforeInvocationId;
}
@JsonProperty("skipSummarization")
@@ -98,19 +94,39 @@ public ConcurrentMap stateDelta() {
return stateDelta;
}
+ @Deprecated // Use stateDelta(), addState() and removeStateByKey() instead.
public void setStateDelta(ConcurrentMap stateDelta) {
this.stateDelta = stateDelta;
}
+ /**
+ * Removes a key from the state delta.
+ *
+ * @param key The key to remove.
+ */
+ public void removeStateByKey(String key) {
+ stateDelta.put(key, State.REMOVED);
+ }
+
@JsonProperty("artifactDelta")
- public ConcurrentMap artifactDelta() {
+ public ConcurrentMap artifactDelta() {
return artifactDelta;
}
- public void setArtifactDelta(ConcurrentMap artifactDelta) {
+ public void setArtifactDelta(ConcurrentMap artifactDelta) {
this.artifactDelta = artifactDelta;
}
+ @JsonProperty("deletedArtifactIds")
+ @JsonInclude(JsonInclude.Include.NON_EMPTY)
+ public Set deletedArtifactIds() {
+ return deletedArtifactIds;
+ }
+
+ public void setDeletedArtifactIds(Set deletedArtifactIds) {
+ this.deletedArtifactIds = deletedArtifactIds;
+ }
+
@JsonProperty("transferToAgent")
public Optional transferToAgent() {
return transferToAgent;
@@ -167,17 +183,28 @@ public void setEndOfAgent(boolean endOfAgent) {
this.endOfAgent = endOfAgent;
}
- @JsonProperty("endInvocation")
+ /**
+ * @deprecated Use {@link #endOfAgent()} instead.
+ */
+ @Deprecated
public Optional endInvocation() {
- return endInvocation;
+ return endOfAgent ? Optional.of(true) : Optional.empty();
}
+ /**
+ * @deprecated Use {@link #setEndOfAgent(boolean)} instead.
+ */
+ @Deprecated
public void setEndInvocation(Optional endInvocation) {
- this.endInvocation = endInvocation;
+ this.endOfAgent = endInvocation.orElse(false);
}
+ /**
+ * @deprecated Use {@link #setEndOfAgent(boolean)} instead.
+ */
+ @Deprecated
public void setEndInvocation(boolean endInvocation) {
- this.endInvocation = Optional.of(endInvocation);
+ this.endOfAgent = endInvocation;
}
@JsonProperty("compaction")
@@ -189,25 +216,6 @@ public void setCompaction(Optional compaction) {
this.compaction = compaction;
}
- @JsonProperty("agentState")
- @JsonInclude(JsonInclude.Include.NON_EMPTY)
- public ConcurrentMap agentState() {
- return agentState;
- }
-
- public void setAgentState(ConcurrentMap agentState) {
- this.agentState = agentState;
- }
-
- @JsonProperty("rewindBeforeInvocationId")
- public Optional rewindBeforeInvocationId() {
- return rewindBeforeInvocationId;
- }
-
- public void setRewindBeforeInvocationId(@Nullable String rewindBeforeInvocationId) {
- this.rewindBeforeInvocationId = Optional.ofNullable(rewindBeforeInvocationId);
- }
-
public static Builder builder() {
return new Builder();
}
@@ -227,15 +235,13 @@ public boolean equals(Object o) {
return Objects.equals(skipSummarization, that.skipSummarization)
&& Objects.equals(stateDelta, that.stateDelta)
&& Objects.equals(artifactDelta, that.artifactDelta)
+ && Objects.equals(deletedArtifactIds, that.deletedArtifactIds)
&& Objects.equals(transferToAgent, that.transferToAgent)
&& Objects.equals(escalate, that.escalate)
&& Objects.equals(requestedAuthConfigs, that.requestedAuthConfigs)
&& Objects.equals(requestedToolConfirmations, that.requestedToolConfirmations)
&& (endOfAgent == that.endOfAgent)
- && Objects.equals(endInvocation, that.endInvocation)
- && Objects.equals(compaction, that.compaction)
- && Objects.equals(agentState, that.agentState)
- && Objects.equals(rewindBeforeInvocationId, that.rewindBeforeInvocationId);
+ && Objects.equals(compaction, that.compaction);
}
@Override
@@ -244,60 +250,52 @@ public int hashCode() {
skipSummarization,
stateDelta,
artifactDelta,
+ deletedArtifactIds,
transferToAgent,
escalate,
requestedAuthConfigs,
requestedToolConfirmations,
endOfAgent,
- endInvocation,
- compaction,
- agentState,
- rewindBeforeInvocationId);
+ compaction);
}
/** Builder for {@link EventActions}. */
public static class Builder {
private Optional skipSummarization;
private ConcurrentMap stateDelta;
- private ConcurrentMap artifactDelta;
+ private ConcurrentMap artifactDelta;
+ private Set deletedArtifactIds;
private Optional transferToAgent;
private Optional escalate;
private ConcurrentMap> requestedAuthConfigs;
private ConcurrentMap requestedToolConfirmations;
private boolean endOfAgent = false;
- private Optional endInvocation;
private Optional compaction;
- private ConcurrentMap agentState;
- private Optional rewindBeforeInvocationId;
public Builder() {
this.skipSummarization = Optional.empty();
this.stateDelta = new ConcurrentHashMap<>();
this.artifactDelta = new ConcurrentHashMap<>();
+ this.deletedArtifactIds = new HashSet<>();
this.transferToAgent = Optional.empty();
this.escalate = Optional.empty();
this.requestedAuthConfigs = new ConcurrentHashMap<>();
this.requestedToolConfirmations = new ConcurrentHashMap<>();
- this.endInvocation = Optional.empty();
this.compaction = Optional.empty();
- this.agentState = new ConcurrentHashMap<>();
- this.rewindBeforeInvocationId = Optional.empty();
}
private Builder(EventActions eventActions) {
this.skipSummarization = eventActions.skipSummarization();
this.stateDelta = new ConcurrentHashMap<>(eventActions.stateDelta());
this.artifactDelta = new ConcurrentHashMap<>(eventActions.artifactDelta());
+ this.deletedArtifactIds = new HashSet<>(eventActions.deletedArtifactIds());
this.transferToAgent = eventActions.transferToAgent();
this.escalate = eventActions.escalate();
this.requestedAuthConfigs = new ConcurrentHashMap<>(eventActions.requestedAuthConfigs());
this.requestedToolConfirmations =
new ConcurrentHashMap<>(eventActions.requestedToolConfirmations());
this.endOfAgent = eventActions.endOfAgent();
- this.endInvocation = eventActions.endInvocation();
this.compaction = eventActions.compaction();
- this.agentState = new ConcurrentHashMap<>(eventActions.agentState());
- this.rewindBeforeInvocationId = eventActions.rewindBeforeInvocationId();
}
@CanIgnoreReturnValue
@@ -316,11 +314,18 @@ public Builder stateDelta(ConcurrentMap value) {
@CanIgnoreReturnValue
@JsonProperty("artifactDelta")
- public Builder artifactDelta(ConcurrentMap value) {
+ public Builder artifactDelta(ConcurrentMap value) {
this.artifactDelta = value;
return this;
}
+ @CanIgnoreReturnValue
+ @JsonProperty("deletedArtifactIds")
+ public Builder deletedArtifactIds(Set value) {
+ this.deletedArtifactIds = value;
+ return this;
+ }
+
@CanIgnoreReturnValue
@JsonProperty("transferToAgent")
public Builder transferToAgent(String agentId) {
@@ -357,10 +362,14 @@ public Builder endOfAgent(boolean endOfAgent) {
return this;
}
+ /**
+ * @deprecated Use {@link #endOfAgent(boolean)} instead.
+ */
@CanIgnoreReturnValue
@JsonProperty("endInvocation")
+ @Deprecated
public Builder endInvocation(boolean endInvocation) {
- this.endInvocation = Optional.of(endInvocation);
+ this.endOfAgent = endInvocation;
return this;
}
@@ -371,34 +380,18 @@ public Builder compaction(EventCompaction value) {
return this;
}
- @CanIgnoreReturnValue
- @JsonProperty("agentState")
- public Builder agentState(ConcurrentMap agentState) {
- this.agentState = agentState;
- 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);
this.stateDelta.putAll(other.stateDelta());
this.artifactDelta.putAll(other.artifactDelta());
+ this.deletedArtifactIds.addAll(other.deletedArtifactIds());
other.transferToAgent().ifPresent(this::transferToAgent);
other.escalate().ifPresent(this::escalate);
this.requestedAuthConfigs.putAll(other.requestedAuthConfigs());
this.requestedToolConfirmations.putAll(other.requestedToolConfirmations());
this.endOfAgent = other.endOfAgent();
- other.endInvocation().ifPresent(this::endInvocation);
other.compaction().ifPresent(this::compaction);
- this.agentState.putAll(other.agentState());
- other.rewindBeforeInvocationId().ifPresent(this::rewindBeforeInvocationId);
return this;
}
diff --git a/core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java b/core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java
index 46b3f1952..6ca49ee62 100644
--- a/core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java
+++ b/core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java
@@ -144,11 +144,15 @@ protected Flowable postprocess(
})
.map(ResponseProcessingResult::updatedResponse);
}
+ Context parentContext = Context.current();
return currentLlmResponse.flatMapPublisher(
- updatedResponse ->
- buildPostprocessingEvents(
- updatedResponse, eventIterables, context, baseEventForLlmResponse, llmRequest));
+ updatedResponse -> {
+ try (Scope scope = parentContext.makeCurrent()) {
+ return buildPostprocessingEvents(
+ updatedResponse, eventIterables, context, baseEventForLlmResponse, llmRequest);
+ }
+ });
}
/**
@@ -175,45 +179,29 @@ private Flowable callLlm(
agent.resolvedModel().model().isPresent()
? agent.resolvedModel().model().get()
: LlmRegistry.getLlm(agent.resolvedModel().modelName().get());
- return Flowable.defer(
- () -> {
- Span llmCallSpan =
- Tracing.getTracer()
- .spanBuilder("call_llm")
- .setParent(Context.current())
- .startSpan();
-
- try (Scope scope = llmCallSpan.makeCurrent()) {
- return llm.generateContent(
- llmRequestBuilder.build(),
- context.runConfig().streamingMode() == StreamingMode.SSE)
- .onErrorResumeNext(
- exception ->
- handleOnModelErrorCallback(
- context,
- llmRequestBuilder,
- eventForCallbackUsage,
- exception)
- .switchIfEmpty(Single.error(exception))
- .toFlowable())
- .doOnNext(
- llmResp -> {
- try (Scope innerScope = llmCallSpan.makeCurrent()) {
- Tracing.traceCallLlm(
- context,
- eventForCallbackUsage.id(),
- llmRequestBuilder.build(),
- llmResp);
- }
- })
- .doOnError(
- error -> {
- llmCallSpan.setStatus(StatusCode.ERROR, error.getMessage());
- llmCallSpan.recordException(error);
- })
- .doFinally(llmCallSpan::end);
- }
+ return llm.generateContent(
+ llmRequestBuilder.build(),
+ context.runConfig().streamingMode() == StreamingMode.SSE)
+ .onErrorResumeNext(
+ exception ->
+ handleOnModelErrorCallback(
+ context, llmRequestBuilder, eventForCallbackUsage, exception)
+ .switchIfEmpty(Single.error(exception))
+ .toFlowable())
+ .doOnNext(
+ llmResp ->
+ Tracing.traceCallLlm(
+ context,
+ eventForCallbackUsage.id(),
+ llmRequestBuilder.build(),
+ llmResp))
+ .doOnError(
+ error -> {
+ Span span = Span.current();
+ span.setStatus(StatusCode.ERROR, error.getMessage());
+ span.recordException(error);
})
+ .compose(Tracing.trace("call_llm"))
.concatMap(
llmResp ->
handleAfterModelCallback(context, llmResp, eventForCallbackUsage)
@@ -230,7 +218,8 @@ private Flowable callLlm(
private Single> handleBeforeModelCallback(
InvocationContext context, LlmRequest.Builder llmRequestBuilder, Event modelResponseEvent) {
Event callbackEvent = modelResponseEvent.toBuilder().build();
- CallbackContext callbackContext = new CallbackContext(context, callbackEvent.actions());
+ CallbackContext callbackContext =
+ new CallbackContext(context, callbackEvent.actions(), callbackEvent.id());
Maybe pluginResult =
context.pluginManager().beforeModelCallback(callbackContext, llmRequestBuilder);
@@ -267,7 +256,8 @@ private Maybe handleOnModelErrorCallback(
Event modelResponseEvent,
Throwable throwable) {
Event callbackEvent = modelResponseEvent.toBuilder().build();
- CallbackContext callbackContext = new CallbackContext(context, callbackEvent.actions());
+ CallbackContext callbackContext =
+ new CallbackContext(context, callbackEvent.actions(), callbackEvent.id());
Exception ex = throwable instanceof Exception e ? e : new Exception(throwable);
Maybe pluginResult =
@@ -301,7 +291,8 @@ private Maybe handleOnModelErrorCallback(
private Single handleAfterModelCallback(
InvocationContext context, LlmResponse llmResponse, Event modelResponseEvent) {
Event callbackEvent = modelResponseEvent.toBuilder().build();
- CallbackContext callbackContext = new CallbackContext(context, callbackEvent.actions());
+ CallbackContext callbackContext =
+ new CallbackContext(context, callbackEvent.actions(), callbackEvent.id());
Maybe pluginResult =
context.pluginManager().afterModelCallback(callbackContext, llmResponse);
@@ -334,73 +325,78 @@ private Single handleAfterModelCallback(
*/
private Flowable runOneStep(InvocationContext context) {
AtomicReference