feat: Add ContextCacheConfig to InvocationContext

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

@@ -55,6 +55,7 @@ public class InvocationContext {
   private final Map<String, Boolean> endOfAgents;
   private final ResumabilityConfig resumabilityConfig;
   @Nullable private final EventsCompactionConfig eventsCompactionConfig;
+  @Nullable private final ContextCacheConfig contextCacheConfig;
   private final InvocationCostManager invocationCostManager;
 
   private Optional<String> branch;
@@ -79,6 +80,7 @@ protected InvocationContext(Builder builder) {
     this.endOfAgents = builder.endOfAgents;
     this.resumabilityConfig = builder.resumabilityConfig;
     this.eventsCompactionConfig = builder.eventsCompactionConfig;
+    this.contextCacheConfig = builder.contextCacheConfig;
     this.invocationCostManager = builder.invocationCostManager;
   }
 
@@ -364,6 +366,11 @@ public Optional<EventsCompactionConfig> eventsCompactionConfig() {
     return Optional.ofNullable(eventsCompactionConfig);
   }
 
+  /** Returns the context cache configuration for the current agent run. */
+  public Optional<ContextCacheConfig> contextCacheConfig() {
+    return Optional.ofNullable(contextCacheConfig);
+  }
+
   /** Returns whether to pause the invocation right after this [event]. */
   public boolean shouldPauseInvocation(Event event) {
     if (!isResumable()) {
@@ -436,6 +443,7 @@ private Builder(InvocationContext context) {
       this.endOfAgents = new ConcurrentHashMap<>(context.endOfAgents);
       this.resumabilityConfig = context.resumabilityConfig;
       this.eventsCompactionConfig = context.eventsCompactionConfig;
+      this.contextCacheConfig = context.contextCacheConfig;
       this.invocationCostManager = context.invocationCostManager;
     }
 
@@ -456,6 +464,7 @@ private Builder(InvocationContext context) {
     private Map<String, Boolean> endOfAgents = new ConcurrentHashMap<>();
     private ResumabilityConfig resumabilityConfig = new ResumabilityConfig();
     @Nullable private EventsCompactionConfig eventsCompactionConfig;
+    @Nullable private ContextCacheConfig contextCacheConfig;
     private InvocationCostManager invocationCostManager = new InvocationCostManager();
 
     /**
@@ -692,6 +701,18 @@ public Builder eventsCompactionConfig(@Nullable EventsCompactionConfig eventsCom
       return this;
     }
 
+    /**
+     * Sets the context cache configuration for the current agent run.
+     *
+     * @param contextCacheConfig the context cache configuration.
+     * @return this builder instance for chaining.
+     */
+    @CanIgnoreReturnValue
+    public Builder contextCacheConfig(@Nullable ContextCacheConfig contextCacheConfig) {
+      this.contextCacheConfig = contextCacheConfig;
+      return this;
+    }
+
     /**
      * Builds the {@link InvocationContext} instance.
      *
@@ -728,6 +749,7 @@ public boolean equals(Object o) {
         && Objects.equals(endOfAgents, that.endOfAgents)
         && Objects.equals(resumabilityConfig, that.resumabilityConfig)
         && Objects.equals(eventsCompactionConfig, that.eventsCompactionConfig)
+        && Objects.equals(contextCacheConfig, that.contextCacheConfig)
         && Objects.equals(invocationCostManager, that.invocationCostManager);
   }
 
@@ -751,6 +773,7 @@ public int hashCode() {
         endOfAgents,
         resumabilityConfig,
         eventsCompactionConfig,
+        contextCacheConfig,
         invocationCostManager);
   }
 }

core/src/main/java/com/google/adk/runner/Runner.java

@@ -18,6 +18,7 @@
 
 import com.google.adk.agents.ActiveStreamingTool;
 import com.google.adk.agents.BaseAgent;
+import com.google.adk.agents.ContextCacheConfig;
 import com.google.adk.agents.InvocationContext;
 import com.google.adk.agents.LiveRequestQueue;
 import com.google.adk.agents.LlmAgent;
@@ -75,6 +76,7 @@ public class Runner {
   private final PluginManager pluginManager;
   private final ResumabilityConfig resumabilityConfig;
   @Nullable private final EventsCompactionConfig eventsCompactionConfig;
+  @Nullable private final ContextCacheConfig contextCacheConfig;
 
   /** Builder for {@link Runner}. */
   public static class Builder {
@@ -138,6 +140,7 @@ public Runner build() {
       List<? extends Plugin> buildPlugins;
       ResumabilityConfig buildResumabilityConfig;
       EventsCompactionConfig buildEventsCompactionConfig;
+      ContextCacheConfig buildContextCacheConfig;
 
       if (this.app != null) {
         if (this.agent != null) {
@@ -154,12 +157,14 @@ public Runner build() {
                 ? this.app.resumabilityConfig()
                 : new ResumabilityConfig();
         buildEventsCompactionConfig = this.app.eventsCompactionConfig();
+        buildContextCacheConfig = this.app.contextCacheConfig();
       } else {
         buildAgent = this.agent;
         buildAppName = this.appName;
         buildPlugins = this.plugins;
         buildResumabilityConfig = new ResumabilityConfig();
         buildEventsCompactionConfig = null;
+        buildContextCacheConfig = null;
       }
 
       if (buildAgent == null) {
@@ -182,7 +187,8 @@ public Runner build() {
           memoryService,
           buildPlugins,
           buildResumabilityConfig,
-          buildEventsCompactionConfig);
+          buildEventsCompactionConfig,
+          buildContextCacheConfig);
     }
   }
 
@@ -257,6 +263,7 @@ public Runner(
         memoryService,
         plugins,
         resumabilityConfig,
+        null,
         null);
   }
 
@@ -274,7 +281,8 @@ protected Runner(
       @Nullable BaseMemoryService memoryService,
       List<? extends Plugin> plugins,
       ResumabilityConfig resumabilityConfig,
-      @Nullable EventsCompactionConfig eventsCompactionConfig) {
+      @Nullable EventsCompactionConfig eventsCompactionConfig,
+      @Nullable ContextCacheConfig contextCacheConfig) {
     this.agent = agent;
     this.appName = appName;
     this.artifactService = artifactService;
@@ -283,6 +291,7 @@ protected Runner(
     this.pluginManager = new PluginManager(plugins);
     this.resumabilityConfig = resumabilityConfig;
     this.eventsCompactionConfig = createEventsCompactionConfig(agent, eventsCompactionConfig);
+    this.contextCacheConfig = contextCacheConfig;
   }
 
   /**
@@ -636,6 +645,7 @@ private InvocationContext.Builder newInvocationContextBuilder(Session session) {
         .session(session)
         .resumabilityConfig(this.resumabilityConfig)
         .eventsCompactionConfig(this.eventsCompactionConfig)
+        .contextCacheConfig(this.contextCacheConfig)
         .agent(this.findAgentToRun(session, rootAgent));
   }