google · copybara-service · Feb 4, 2026
diff --git a/core/src/main/java/com/google/adk/agents/InvocationContext.java b/core/src/main/java/com/google/adk/agents/InvocationContext.java
@@ -56,6 +56,7 @@ public class InvocationContext {
   private final ResumabilityConfig resumabilityConfig;
   @Nullable private final EventsCompactionConfig eventsCompactionConfig;
   private final InvocationCostManager invocationCostManager;
+  private final Map<String, Object> callbackContextData;
 
   private Optional<String> branch;
   private BaseAgent agent;
@@ -80,6 +81,7 @@ protected InvocationContext(Builder builder) {
     this.resumabilityConfig = builder.resumabilityConfig;
     this.eventsCompactionConfig = builder.eventsCompactionConfig;
     this.invocationCostManager = builder.invocationCostManager;
+    this.callbackContextData = builder.callbackContextData;
   }
 
   /**
@@ -306,6 +308,14 @@ public RunConfig runConfig() {
     return runConfig;
   }
 
+  /**
+   * 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<String, Object> callbackContextData() {
+    return callbackContextData;
+  }
+
   /** Returns agent-specific state saved within this invocation. */
   public Map<String, BaseAgentState> agentStates() {
     return agentStates;
@@ -437,6 +447,7 @@ private Builder(InvocationContext context) {
       this.resumabilityConfig = context.resumabilityConfig;
       this.eventsCompactionConfig = context.eventsCompactionConfig;
       this.invocationCostManager = context.invocationCostManager;
+      this.callbackContextData = context.callbackContextData;
     }
 
     private BaseSessionService sessionService;
@@ -457,6 +468,7 @@ private Builder(InvocationContext context) {
     private ResumabilityConfig resumabilityConfig = new ResumabilityConfig();
     @Nullable private EventsCompactionConfig eventsCompactionConfig;
     private InvocationCostManager invocationCostManager = new InvocationCostManager();
+    private Map<String, Object> callbackContextData = new ConcurrentHashMap<>();
 
     /**
      * Sets the session service for managing session state.
@@ -692,6 +704,18 @@ public Builder eventsCompactionConfig(@Nullable EventsCompactionConfig eventsCom
       return this;
     }
 
+    /**
+     * Sets the callback context data for the invocation.
+     *
+     * @param callbackContextData the callback context data.
+     * @return this builder instance for chaining.
+     */
+    @CanIgnoreReturnValue
+    public Builder callbackContextData(Map<String, Object> callbackContextData) {
+      this.callbackContextData = callbackContextData;
+      return this;
+    }
+
     /**
      * Builds the {@link InvocationContext} instance.
      *
@@ -728,7 +752,8 @@ public boolean equals(Object o) {
         && Objects.equals(endOfAgents, that.endOfAgents)
         && Objects.equals(resumabilityConfig, that.resumabilityConfig)
         && Objects.equals(eventsCompactionConfig, that.eventsCompactionConfig)
-        && Objects.equals(invocationCostManager, that.invocationCostManager);
+        && Objects.equals(invocationCostManager, that.invocationCostManager)
+        && Objects.equals(callbackContextData, that.callbackContextData);
   }
 
   @Override
@@ -751,6 +776,7 @@ public int hashCode() {
         endOfAgents,
         resumabilityConfig,
         eventsCompactionConfig,
-        invocationCostManager);
+        invocationCostManager,
+        callbackContextData);
   }
 }
diff --git a/core/src/test/java/com/google/adk/agents/InvocationContextTest.java b/core/src/test/java/com/google/adk/agents/InvocationContextTest.java
@@ -27,13 +27,15 @@
 import com.google.adk.sessions.BaseSessionService;
 import com.google.adk.sessions.Session;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.genai.types.Content;
 import com.google.genai.types.FunctionCall;
 import com.google.genai.types.Part;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Optional;
+import java.util.concurrent.ConcurrentHashMap;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -178,6 +180,25 @@ public void testCopyOf() {
     assertThat(copiedContext.endInvocation()).isEqualTo(originalContext.endInvocation());
     assertThat(copiedContext.activeStreamingTools())
         .isEqualTo(originalContext.activeStreamingTools());
+    assertThat(copiedContext.callbackContextData())
+        .isSameInstanceAs(originalContext.callbackContextData());
+  }
+
+  @Test
+  public void testBuildWithCallbackContextData() {
+    Map<String, Object> data = new ConcurrentHashMap<>();
+    data.put("key", "value");
+    InvocationContext context =
+        InvocationContext.builder()
+            .sessionService(mockSessionService)
+            .artifactService(mockArtifactService)
+            .agent(mockAgent)
+            .session(session)
+            .callbackContextData(data)
+            .build();
+
+    assertThat(context.callbackContextData()).isEqualTo(data);
+    assertThat(context.callbackContextData()).isSameInstanceAs(data);
   }
 
   @Test
@@ -404,6 +425,22 @@ public void testEquals_differentValues() {
     assertThat(context.equals(contextWithDiffAgent)).isFalse();
     assertThat(context.equals(contextWithUserContentEmpty)).isFalse();
     assertThat(context.equals(contextWithLiveQueuePresent)).isFalse();
+
+    InvocationContext contextWithDiffCallbackContextData =
+        InvocationContext.builder()
+            .sessionService(mockSessionService)
+            .artifactService(mockArtifactService)
+            .memoryService(mockMemoryService)
+            .pluginManager(pluginManager)
+            .invocationId(testInvocationId)
+            .agent(mockAgent)
+            .session(session)
+            .userContent(userContent)
+            .runConfig(runConfig)
+            .endInvocation(false)
+            .callbackContextData(ImmutableMap.of("key", "value"))
+            .build();
+    assertThat(context.equals(contextWithDiffCallbackContextData)).isFalse();
   }
 
   @Test
@@ -453,6 +490,22 @@ public void testHashCode_differentValues() {
 
     assertThat(context).isNotEqualTo(contextWithDiffSessionService);
     assertThat(context).isNotEqualTo(contextWithDiffInvocationId);
+
+    InvocationContext contextWithDiffCallbackContextData =
+        InvocationContext.builder()
+            .sessionService(mockSessionService)
+            .artifactService(mockArtifactService)
+            .memoryService(mockMemoryService)
+            .pluginManager(pluginManager)
+            .invocationId(testInvocationId)
+            .agent(mockAgent)
+            .session(session)
+            .userContent(userContent)
+            .runConfig(runConfig)
+            .endInvocation(false)
+            .callbackContextData(ImmutableMap.of("key", "value"))
+            .build();
+    assertThat(context.hashCode()).isNotEqualTo(contextWithDiffCallbackContextData.hashCode());
   }
 
   @Test

diff --git a/core/src/test/java/com/google/adk/flows/llmflows/BaseLlmFlowTest.java b/core/src/test/java/com/google/adk/flows/llmflows/BaseLlmFlowTest.java
@@ -25,6 +25,7 @@
 import static com.google.common.collect.Iterables.getOnlyElement;
 import static com.google.common.truth.Truth.assertThat;
 
+import com.google.adk.agents.Callbacks;
 import com.google.adk.agents.InvocationContext;
 import com.google.adk.events.Event;
 import com.google.adk.flows.llmflows.RequestProcessor.RequestProcessingResult;
@@ -42,6 +43,7 @@
 import com.google.genai.types.GenerateContentResponseUsageMetadata;
 import com.google.genai.types.Part;
 import io.reactivex.rxjava3.core.Flowable;
+import io.reactivex.rxjava3.core.Maybe;
 import io.reactivex.rxjava3.core.Single;
 import java.util.List;
 import java.util.Map;
@@ -414,6 +416,82 @@ public void run_requestProcessorsAreCalledExactlyOnce() {
     assertThat(processor2CallCount.get()).isEqualTo(1);
   }
 
+  @Test
+  public void run_sharingcallbackContextDataBetweenCallbacks() {
+    Content content = Content.fromParts(Part.fromText("LLM response"));
+    TestLlm testLlm = createTestLlm(createLlmResponse(content));
+
+    Callbacks.BeforeModelCallback beforeCallback =
+        (ctx, req) -> {
+          ctx.invocationContext().callbackContextData().put("key", "value_from_before");
+          return Maybe.empty();
+        };
+
+    Callbacks.AfterModelCallback afterCallback =
+        (ctx, resp) -> {
+          String value = (String) ctx.invocationContext().callbackContextData().get("key");
+          LlmResponse modifiedResp =
+              resp.toBuilder().content(Content.fromParts(Part.fromText("Saw: " + value))).build();
+          return Maybe.just(modifiedResp);
+        };
+
+    InvocationContext invocationContext =
+        createInvocationContext(
+            createTestAgentBuilder(testLlm)
+                .beforeModelCallback(beforeCallback)
+                .afterModelCallback(afterCallback)
+                .build());
+
+    BaseLlmFlow baseLlmFlow = createBaseLlmFlowWithoutProcessors();
+
+    List<Event> events = baseLlmFlow.run(invocationContext).toList().blockingGet();
+
+    assertThat(events).hasSize(1);
+    assertThat(events.get(0).stringifyContent()).isEqualTo("Saw: value_from_before");
+  }
+
+  @Test
+  public void run_sharingcallbackContextDataAcrossContextCopies() {
+    Content content = Content.fromParts(Part.fromText("LLM response"));
+    TestLlm testLlm = createTestLlm(createLlmResponse(content));
+
+    Callbacks.BeforeModelCallback beforeCallback =
+        (ctx, req) -> {
+          ctx.invocationContext().callbackContextData().put("key", "value_from_before");
+          return Maybe.empty();
+        };
+
+    Callbacks.AfterModelCallback afterCallback =
+        (ctx, resp) -> {
+          String value = (String) ctx.invocationContext().callbackContextData().get("key");
+          LlmResponse modifiedResp =
+              resp.toBuilder().content(Content.fromParts(Part.fromText("Saw: " + value))).build();
+          return Maybe.just(modifiedResp);
+        };
+
+    InvocationContext invocationContext =
+        createInvocationContext(
+            createTestAgentBuilder(testLlm)
+                .beforeModelCallback(beforeCallback)
+                .afterModelCallback(afterCallback)
+                .build());
+
+    BaseLlmFlow baseLlmFlow =
+        new BaseLlmFlow(ImmutableList.of(), ImmutableList.of()) {
+          @Override
+          public Flowable<Event> run(InvocationContext context) {
+            // Force a context copy
+            InvocationContext copiedContext = context.toBuilder().build();
+            return super.run(copiedContext);
+          }
+        };
+
+    List<Event> events = baseLlmFlow.run(invocationContext).toList().blockingGet();
+
+    assertThat(events).hasSize(1);
+    assertThat(events.get(0).stringifyContent()).isEqualTo("Saw: value_from_before");
+  }
+
   private static BaseLlmFlow createBaseLlmFlowWithoutProcessors() {
     return createBaseLlmFlow(ImmutableList.of(), ImmutableList.of());
   }