(feat) : implement cancellation support with CancellationToken and CancellationTokenSource

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

@@ -26,7 +26,7 @@
 import com.google.adk.plugins.Plugin;
 import com.google.adk.telemetry.Instrumentation;
 import com.google.adk.telemetry.Instrumentation.AgentInvocation;
-import com.google.adk.utils.AgentEnums.AgentOrigin;
+import com.google.adk.utils.AgentOrigin;
 import com.google.common.collect.ImmutableList;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import com.google.errorprone.annotations.DoNotCall;
@@ -331,16 +331,26 @@ private Flowable<Event> run(
         },
         agentInvocation -> {
           InvocationContext invocationContext = agentInvocation.getCtx();
+          if (invocationContext.isCancellationRequested()) {
+            return Flowable.empty();
+          }
           Flowable<Event> mainAndAfterEvents =
-              Flowable.defer(() -> runImplementation.apply(invocationContext))
+              Flowable.defer(
+                      () ->
+                          invocationContext.isCancellationRequested()
+                              ? Flowable.empty()
+                              : runImplementation.apply(invocationContext))
                   .concatWith(
                       Flowable.defer(
                           () ->
-                              callCallback(
-                                      afterCallbacksToFunctions(
-                                          invocationContext.pluginManager(), afterAgentCallback),
-                                      invocationContext)
-                                  .toFlowable()));
+                              invocationContext.isCancellationRequested()
+                                  ? Flowable.empty()
+                                  : callCallback(
+                                          afterCallbacksToFunctions(
+                                              invocationContext.pluginManager(),
+                                              afterAgentCallback),
+                                          invocationContext)
+                                      .toFlowable()));
 
           return callCallback(
                   beforeCallbacksToFunctions(

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

@@ -0,0 +1,36 @@
+/*
+ * 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
+ *
+ *     https://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;
+
+/** A thread-safe, cooperative cancellation signal for an agent invocation. */
+@FunctionalInterface
+public interface CancellationToken {
+  /** Returns whether cancellation has been requested. */
+  boolean isCancellationRequested();
+
+  /** Returns a token that is never cancelled. */
+  static CancellationToken none() {
+    return NeverCancelledHolder.INSTANCE;
+  }
+
+  /** Holder for the shared no-op token. */
+  final class NeverCancelledHolder {
+    private static final CancellationToken INSTANCE = () -> false;
+
+    private NeverCancelledHolder() {}
+  }
+}

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

@@ -0,0 +1,46 @@
+/*
+ * 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
+ *
+ *     https://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.util.concurrent.atomic.AtomicBoolean;
+
+/** Owns a {@link CancellationToken} and can request cancellation for it. */
+public final class CancellationTokenSource implements AutoCloseable {
+  private final AtomicBoolean cancellationRequested = new AtomicBoolean();
+  private final CancellationToken token = cancellationRequested::get;
+
+  /** Returns the token controlled by this source. */
+  public CancellationToken token() {
+    return token;
+  }
+
+  /** Requests cancellation. Calling this method more than once is safe. */
+  public void cancel() {
+    cancellationRequested.set(true);
+  }
+
+  /** Returns whether cancellation has been requested. */
+  public boolean isCancellationRequested() {
+    return cancellationRequested.get();
+  }
+
+  /** Requests cancellation. */
+  @Override
+  public void close() {
+    cancel();
+  }
+}

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

@@ -166,6 +166,16 @@ public RunConfig runConfig() {
     return runConfig;
   }
 
+  /** Returns the cancellation token for this invocation. */
+  public CancellationToken cancellationToken() {
+    return runConfig.cancellationToken();
+  }
+
+  /** Returns whether cancellation has been requested for this invocation. */
+  public boolean isCancellationRequested() {
+    return cancellationToken().isCancellationRequested();
+  }
+
   /**
    * 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).

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

@@ -117,6 +117,9 @@ public final boolean groupFunctionResponsesInHistory() {
 
   public abstract ImmutableMap<String, Object> customMetadata();
 
+  /** Token used to cooperatively cancel this run. */
+  public abstract CancellationToken cancellationToken();
+
   public abstract Builder toBuilder();
 
   public static Builder builder() {
@@ -128,6 +131,7 @@ public static Builder builder() {
         .toolExecutionMode(ToolExecutionMode.NONE)
         .maxLlmCalls(500)
         .autoCreateSession(false)
+        .cancellationToken(CancellationToken.none())
         .customMetadata(ImmutableMap.of());
   }
 
@@ -144,6 +148,7 @@ public static Builder builder(RunConfig runConfig) {
         .outputAudioTranscription(runConfig.outputAudioTranscription())
         .inputAudioTranscription(runConfig.inputAudioTranscription())
         .autoCreateSession(runConfig.autoCreateSession())
+        .cancellationToken(runConfig.cancellationToken())
         .groupFunctionResponsesInHistoryOverride(
             runConfig.groupFunctionResponsesInHistoryOverride())
         .customMetadata(runConfig.customMetadata());
@@ -244,6 +249,9 @@ public final Builder setAutoCreateSession(boolean autoCreateSession) {
     @CanIgnoreReturnValue
     public abstract Builder customMetadata(Map<String, Object> customMetadata);
 
+    @CanIgnoreReturnValue
+    public abstract Builder cancellationToken(CancellationToken cancellationToken);
+
     /**
      * Sets the three-state grouping override.
      *

core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java

@@ -106,7 +106,13 @@ private Flowable<Event> preprocess(
     return Flowable.fromIterable(allProcessors)
         .concatMap(
             processor ->
-                Single.defer(() -> processor.processRequest(context, llmRequestRef.get()))
+                Single.defer(
+                        () ->
+                            context.isCancellationRequested()
+                                ? Single.just(
+                                    RequestProcessingResult.create(
+                                        llmRequestRef.get(), ImmutableList.of()))
+                                : processor.processRequest(context, llmRequestRef.get()))
                     .compose(Tracing.withContext(currentContext))
                     .doOnSuccess(result -> llmRequestRef.set(result.updatedRequest()))
                     .flattenAsFlowable(
@@ -182,7 +188,12 @@ protected Flowable<Event> postprocess(
     for (ResponseProcessor processor : responseProcessors) {
       currentLlmResponse =
           currentLlmResponse
-              .flatMap(response -> processor.processResponse(context, response))
+              .flatMap(
+                  response ->
+                      context.isCancellationRequested()
+                          ? Single.just(
+                              ResponseProcessingResult.create(response, ImmutableList.of()))
+                          : processor.processResponse(context, response))
               .doOnSuccess(
                   result -> {
                     if (result.events() != null) {
@@ -220,6 +231,9 @@ private Flowable<Event> callLlm(
 
     return Flowable.defer(
         () -> {
+          if (context.isCancellationRequested()) {
+            return Flowable.empty();
+          }
           Span span =
               Tracing.getTracer().spanBuilder("call_llm").setParent(spanContext).startSpan();
           Context callLlmContext = spanContext.with(span);
@@ -249,6 +263,9 @@ private Flowable<Event> callLlm(
                           .switchIfEmpty(
                               Flowable.defer(
                                   () -> {
+                                    if (context.isCancellationRequested()) {
+                                      return Flowable.empty();
+                                    }
                                     LlmAgent agent = (LlmAgent) context.agent();
                                     BaseLlm llm =
                                         agent.resolvedModel().model().isPresent()
@@ -277,9 +294,11 @@ private Flowable<Event> callLlm(
                                             })
                                         .concatMap(
                                             llmResp ->
-                                                handleAfterModelCallback(
-                                                        context, llmResp, eventForCallbackUsage)
-                                                    .toFlowable())
+                                                context.isCancellationRequested()
+                                                    ? Flowable.empty()
+                                                    : handleAfterModelCallback(
+                                                            context, llmResp, eventForCallbackUsage)
+                                                        .toFlowable())
                                         .flatMap(
                                             llmResp ->
                                                 postprocess(
@@ -309,6 +328,9 @@ private Flowable<Event> callLlm(
    */
   private Maybe<LlmResponse> handleBeforeModelCallback(
       InvocationContext context, LlmRequest.Builder llmRequestBuilder, Event modelResponseEvent) {
+    if (context.isCancellationRequested()) {
+      return Maybe.empty();
+    }
     Context currentContext = Context.current();
     Event callbackEvent = modelResponseEvent.toBuilder().build();
     CallbackContext callbackContext =
@@ -437,7 +459,7 @@ private Flowable<Event> runOneStep(Context spanContext, InvocationContext contex
                   Flowable.defer(
                       () -> {
                         LlmRequest llmRequestAfterPreprocess = llmRequestRef.get();
-                        if (context.endInvocation()) {
+                        if (context.endInvocation() || context.isCancellationRequested()) {
                           logger.debug("End invocation requested during preprocessing.");
                           return Flowable.empty();
                         }
@@ -528,6 +550,7 @@ private Flowable<Event> run(
             .flatMapPublisher(
                 eventList -> {
                   if (eventList.isEmpty()
+                      || invocationContext.isCancellationRequested()
                       || Iterables.getLast(eventList).finalResponse()
                       || Iterables.getLast(eventList).actions().endInvocation().orElse(false)) {
                     logger.debug(
@@ -571,7 +594,8 @@ public Flowable<Event> runLive(InvocationContext invocationContext) {
         Flowable.defer(
             () -> {
               LlmRequest llmRequestAfterPreprocess = llmRequestRef.get();
-              if (invocationContext.endInvocation()) {
+              if (invocationContext.endInvocation()
+                  || invocationContext.isCancellationRequested()) {
                 return Flowable.empty();
               }
 
@@ -664,6 +688,7 @@ public void onError(Throwable e) {
               Flowable<Event> receiveFlow =
                   connection
                       .receive()
+                      .takeWhile(unused -> !invocationContext.isCancellationRequested())
                       .flatMap(
                           llmResponse -> {
                             Event baseEventForThisLlmResponse =
@@ -677,6 +702,9 @@ public void onError(Throwable e) {
                           })
                       .flatMap(
                           event -> {
+                            if (invocationContext.isCancellationRequested()) {
+                              return Flowable.empty();
+                            }
                             Flowable<Event> events = Flowable.just(event);
                             if (event.actions().transferToAgent().isPresent()) {
                               BaseAgent rootAgent = invocationContext.agent().rootAgent();
@@ -712,6 +740,11 @@ public void onError(Throwable e) {
                               sendTask.dispose();
                               connection.close();
                             }
+                          })
+                      .doFinally(
+                          () -> {
+                            sendTask.dispose();
+                            connection.close();
                           });
 
               return Tracing.traceFlowable(

core/src/main/java/com/google/adk/flows/llmflows/Functions.java

@@ -292,6 +292,9 @@ private static Function<FunctionCall, Maybe<Event>> getFunctionCallMapper(
     return functionCall ->
         Maybe.defer(
                 () -> {
+                  if (invocationContext.isCancellationRequested()) {
+                    return Maybe.empty();
+                  }
                   BaseTool tool = tools.get(functionCall.name().get());
                   ToolContext toolContext =
                       ToolContext.builder(invocationContext)
@@ -308,14 +311,16 @@ private static Function<FunctionCall, Maybe<Event>> getFunctionCallMapper(
                           .switchIfEmpty(
                               Maybe.defer(
                                       () ->
-                                          isLive
-                                              ? processFunctionLive(
-                                                  invocationContext,
-                                                  tool,
-                                                  toolContext,
-                                                  functionCall,
-                                                  functionArgs)
-                                              : callTool(tool, functionArgs, toolContext))
+                                          invocationContext.isCancellationRequested()
+                                              ? Maybe.empty()
+                                              : isLive
+                                                  ? processFunctionLive(
+                                                      invocationContext,
+                                                      tool,
+                                                      toolContext,
+                                                      functionCall,
+                                                      functionArgs)
+                                                  : callTool(tool, functionArgs, toolContext))
                                   .compose(Tracing.withContext(parentContext)));
 
                   return postProcessFunctionResult(
@@ -540,6 +545,9 @@ private static Maybe<Event> processFunctionResult(
                   .defaultIfEmpty(Optional.ofNullable(initialFunctionResult))
                   .flatMapMaybe(
                       finalOptionalResult -> {
+                        if (invocationContext.isCancellationRequested()) {
+                          return Maybe.empty();
+                        }
                         Map<String, Object> finalFunctionResult = finalOptionalResult.orElse(null);
                         boolean hasNoResult =
                             finalFunctionResult == null || finalFunctionResult.isEmpty();
@@ -599,6 +607,9 @@ private static Maybe<Map<String, Object>> maybeInvokeBeforeToolCall(
       BaseTool tool,
       Map<String, Object> functionArgs,
       ToolContext toolContext) {
+    if (invocationContext.isCancellationRequested()) {
+      return Maybe.empty();
+    }
     if (invocationContext.agent() instanceof LlmAgent) {
       LlmAgent agent = (LlmAgent) invocationContext.agent();
 
@@ -671,6 +682,9 @@ private static Maybe<Map<String, Object>> maybeInvokeAfterToolCall(
       Map<String, Object> functionArgs,
       ToolContext toolContext,
       Map<String, Object> functionResult) {
+    if (invocationContext.isCancellationRequested()) {
+      return Maybe.empty();
+    }
     if (invocationContext.agent() instanceof LlmAgent) {
       LlmAgent agent = (LlmAgent) invocationContext.agent();
 

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

@@ -524,6 +524,9 @@ protected Flowable<Event> runAsyncImpl(
     Preconditions.checkNotNull(runConfig, "runConfig cannot be null");
     return Flowable.defer(
             () -> {
+              if (runConfig.cancellationToken().isCancellationRequested()) {
+                return Flowable.empty();
+              }
               Context capturedContext = Context.current();
               BaseAgent rootAgent = this.agent;
               String invocationId = InvocationContext.newInvocationContextId();
@@ -778,6 +781,9 @@ protected Flowable<Event> runLiveImpl(
       Session session, @Nullable LiveRequestQueue liveRequestQueue, RunConfig runConfig) {
     return Flowable.defer(
         () -> {
+          if (runConfig.cancellationToken().isCancellationRequested()) {
+            return Flowable.empty();
+          }
           Context capturedContext = Context.current();
           InvocationContext invocationContext =
               newInvocationContextForLive(session, liveRequestQueue, runConfig);