feat: add eventId in CallbackContext and ToolContext

google-genai-bot · copybara-github · commit ac05fde31ec6 · 2026-02-05T12:00:26.000-08:00
This allows callbacks and tools to know which event triggered them.  This is useful in telemetry and tracing.

PiperOrigin-RevId: 866050613
diff --git a/core/src/main/java/com/google/adk/agents/CallbackContext.java b/core/src/main/java/com/google/adk/agents/CallbackContext.java
@@ -31,6 +31,7 @@ public class CallbackContext extends ReadonlyContext {
 
   protected EventActions eventActions;
   private final State state;
+  private final String eventId;
 
   /**
    * Initializes callback context.
@@ -39,9 +40,22 @@ public class CallbackContext extends ReadonlyContext {
    * @param eventActions Callback event actions.
    */
   public CallbackContext(InvocationContext invocationContext, EventActions eventActions) {
+    this(invocationContext, eventActions, null);
+  }
+
+  /**
+   * Initializes callback context.
+   *
+   * @param invocationContext Current invocation context.
+   * @param eventActions Callback event actions.
+   * @param eventId The ID of the event associated with this context.
+   */
+  public CallbackContext(
+      InvocationContext invocationContext, EventActions eventActions, String eventId) {
     super(invocationContext);
     this.eventActions = eventActions != null ? eventActions : EventActions.builder().build();
     this.state = new State(invocationContext.session().state(), this.eventActions.stateDelta());
+    this.eventId = eventId;
   }
 
   /** Returns the delta-aware state of the current callback. */
@@ -55,6 +69,11 @@ public EventActions eventActions() {
     return eventActions;
   }
 
+  /** Returns the ID of the event associated with this context. */
+  public String eventId() {
+    return eventId;
+  }
+
   /**
    * Lists the filenames of the artifacts attached to the current session.
    *
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
@@ -237,7 +237,8 @@ private Flowable<LlmResponse> callLlm(
   private Single<Optional<LlmResponse>> 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<LlmResponse> pluginResult =
         context.pluginManager().beforeModelCallback(callbackContext, llmRequestBuilder);
@@ -274,7 +275,8 @@ private Maybe<LlmResponse> 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<LlmResponse> pluginResult =
@@ -308,7 +310,8 @@ private Maybe<LlmResponse> handleOnModelErrorCallback(
   private Single<LlmResponse> 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<LlmResponse> pluginResult =
         context.pluginManager().afterModelCallback(callbackContext, llmResponse);
diff --git a/core/src/main/java/com/google/adk/tools/ToolContext.java b/core/src/main/java/com/google/adk/tools/ToolContext.java
@@ -35,8 +35,9 @@ private ToolContext(
       InvocationContext invocationContext,
       EventActions eventActions,
       Optional<String> functionCallId,
-      Optional<ToolConfirmation> toolConfirmation) {
-    super(invocationContext, eventActions);
+      Optional<ToolConfirmation> toolConfirmation,
+      @Nullable String eventId) {
+    super(invocationContext, eventActions, eventId);
     this.functionCallId = functionCallId;
     this.toolConfirmation = toolConfirmation;
   }
@@ -125,7 +126,8 @@ public Builder toBuilder() {
     return new Builder(invocationContext)
         .actions(eventActions)
         .functionCallId(functionCallId.orElse(null))
-        .toolConfirmation(toolConfirmation.orElse(null));
+        .toolConfirmation(toolConfirmation.orElse(null))
+        .eventId(eventId());
   }
 
   @Override
@@ -148,6 +150,7 @@ public static final class Builder {
     private EventActions eventActions = EventActions.builder().build(); // Default empty actions
     private Optional<String> functionCallId = Optional.empty();
     private Optional<ToolConfirmation> toolConfirmation = Optional.empty();
+    private String eventId;
 
     private Builder(InvocationContext invocationContext) {
       this.invocationContext = invocationContext;
@@ -171,8 +174,15 @@ public Builder toolConfirmation(ToolConfirmation toolConfirmation) {
       return this;
     }
 
+    @CanIgnoreReturnValue
+    public Builder eventId(String eventId) {
+      this.eventId = eventId;
+      return this;
+    }
+
     public ToolContext build() {
-      return new ToolContext(invocationContext, eventActions, functionCallId, toolConfirmation);
+      return new ToolContext(
+          invocationContext, eventActions, functionCallId, toolConfirmation, eventId);
     }
   }
 }
