diff --git a/core/src/main/java/com/google/adk/agents/BaseAgent.java b/core/src/main/java/com/google/adk/agents/BaseAgent.java index fc1f0f31e..db89c4c25 100644 --- a/core/src/main/java/com/google/adk/agents/BaseAgent.java +++ b/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 run( }, agentInvocation -> { InvocationContext invocationContext = agentInvocation.getCtx(); + if (invocationContext.isCancellationRequested()) { + return Flowable.empty(); + } Flowable 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( diff --git a/core/src/main/java/com/google/adk/agents/CancellationToken.java b/core/src/main/java/com/google/adk/agents/CancellationToken.java new file mode 100644 index 000000000..fb932e482 --- /dev/null +++ b/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() {} + } +} diff --git a/core/src/main/java/com/google/adk/agents/CancellationTokenSource.java b/core/src/main/java/com/google/adk/agents/CancellationTokenSource.java new file mode 100644 index 000000000..e07c9a2f4 --- /dev/null +++ b/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(); + } +} 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 456758b95..5f13ee8e7 100644 --- a/core/src/main/java/com/google/adk/agents/InvocationContext.java +++ b/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). 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 bd20b6183..a3be707d9 100644 --- a/core/src/main/java/com/google/adk/agents/RunConfig.java +++ b/core/src/main/java/com/google/adk/agents/RunConfig.java @@ -117,6 +117,9 @@ public final boolean groupFunctionResponsesInHistory() { public abstract ImmutableMap 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 customMetadata); + @CanIgnoreReturnValue + public abstract Builder cancellationToken(CancellationToken cancellationToken); + /** * Sets the three-state grouping override. * 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 91cc225f2..28d535ee6 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 @@ -106,7 +106,13 @@ private Flowable 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 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 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 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 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 callLlm( */ private Maybe 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 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 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 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 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 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( diff --git a/core/src/main/java/com/google/adk/flows/llmflows/Functions.java b/core/src/main/java/com/google/adk/flows/llmflows/Functions.java index 3f3b8ef86..71e343b7c 100644 --- a/core/src/main/java/com/google/adk/flows/llmflows/Functions.java +++ b/core/src/main/java/com/google/adk/flows/llmflows/Functions.java @@ -292,6 +292,9 @@ private static Function> 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> 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 processFunctionResult( .defaultIfEmpty(Optional.ofNullable(initialFunctionResult)) .flatMapMaybe( finalOptionalResult -> { + if (invocationContext.isCancellationRequested()) { + return Maybe.empty(); + } Map finalFunctionResult = finalOptionalResult.orElse(null); boolean hasNoResult = finalFunctionResult == null || finalFunctionResult.isEmpty(); @@ -599,6 +607,9 @@ private static Maybe> maybeInvokeBeforeToolCall( BaseTool tool, Map 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> maybeInvokeAfterToolCall( Map functionArgs, ToolContext toolContext, Map functionResult) { + if (invocationContext.isCancellationRequested()) { + return Maybe.empty(); + } if (invocationContext.agent() instanceof LlmAgent) { LlmAgent agent = (LlmAgent) invocationContext.agent(); diff --git a/core/src/main/java/com/google/adk/runner/Runner.java b/core/src/main/java/com/google/adk/runner/Runner.java index 48eb9fad6..e7b7d8d02 100644 --- a/core/src/main/java/com/google/adk/runner/Runner.java +++ b/core/src/main/java/com/google/adk/runner/Runner.java @@ -524,6 +524,9 @@ protected Flowable 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 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); diff --git a/core/src/test/java/com/google/adk/agents/CancellationTokenSourceTest.java b/core/src/test/java/com/google/adk/agents/CancellationTokenSourceTest.java new file mode 100644 index 000000000..553fa6a28 --- /dev/null +++ b/core/src/test/java/com/google/adk/agents/CancellationTokenSourceTest.java @@ -0,0 +1,49 @@ +/* + * 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 static com.google.common.truth.Truth.assertThat; + +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public final class CancellationTokenSourceTest { + + @Test + public void cancel_isVisibleThroughTokenAndIdempotent() { + CancellationTokenSource source = new CancellationTokenSource(); + + assertThat(source.token().isCancellationRequested()).isFalse(); + + source.cancel(); + source.cancel(); + + assertThat(source.isCancellationRequested()).isTrue(); + assertThat(source.token().isCancellationRequested()).isTrue(); + } + + @Test + public void close_requestsCancellation() { + CancellationTokenSource source = new CancellationTokenSource(); + + source.close(); + + assertThat(source.token().isCancellationRequested()).isTrue(); + } +} diff --git a/core/src/test/java/com/google/adk/agents/RunConfigTest.java b/core/src/test/java/com/google/adk/agents/RunConfigTest.java index fc6b9083f..5104a50bf 100644 --- a/core/src/test/java/com/google/adk/agents/RunConfigTest.java +++ b/core/src/test/java/com/google/adk/agents/RunConfigTest.java @@ -74,6 +74,18 @@ public void testBuilderDefaults() { assertThat(runConfig.autoCreateSession()).isFalse(); assertThat(runConfig.groupFunctionResponsesInHistoryOverride()).isEmpty(); assertThat(runConfig.groupFunctionResponsesInHistory()).isFalse(); + assertThat(runConfig.cancellationToken().isCancellationRequested()).isFalse(); + } + + @Test + public void cancellationToken_isPropagatedByCopyBuilder() { + CancellationTokenSource source = new CancellationTokenSource(); + RunConfig runConfig = RunConfig.builder().cancellationToken(source.token()).build(); + + RunConfig copy = RunConfig.builder(runConfig).build(); + source.cancel(); + + assertThat(copy.cancellationToken().isCancellationRequested()).isTrue(); } @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 index 1761871e6..3f5317269 100644 --- a/core/src/test/java/com/google/adk/flows/llmflows/BaseLlmFlowTest.java +++ b/core/src/test/java/com/google/adk/flows/llmflows/BaseLlmFlowTest.java @@ -28,6 +28,7 @@ import static org.junit.Assert.assertThrows; import com.google.adk.agents.Callbacks; +import com.google.adk.agents.CancellationTokenSource; import com.google.adk.agents.InvocationContext; import com.google.adk.agents.LlmAgent; import com.google.adk.agents.ReadonlyContext; @@ -72,6 +73,58 @@ @RunWith(JUnit4.class) public final class BaseLlmFlowTest { + @Test + public void run_cancelledBeforeSubscription_completesWithoutCallingModel() { + AtomicInteger modelCalls = new AtomicInteger(); + TestLlm testLlm = + createTestLlm( + () -> { + modelCalls.incrementAndGet(); + return Flowable.just( + createLlmResponse(Content.fromParts(Part.fromText("unreachable")))); + }); + CancellationTokenSource cancellation = new CancellationTokenSource(); + InvocationContext invocationContext = + createInvocationContext( + createTestAgent(testLlm), + RunConfig.builder().cancellationToken(cancellation.token()).build()); + BaseLlmFlow baseLlmFlow = createBaseLlmFlowWithoutProcessors(); + cancellation.cancel(); + + baseLlmFlow.run(invocationContext).test().assertComplete().assertNoErrors().assertNoValues(); + + assertThat(modelCalls.get()).isEqualTo(0); + } + + @Test + public void run_cancelledBeforeToolExecution_preservesModelEventAndSkipsTool() { + Content functionCall = + Content.fromParts(Part.fromFunctionCall("counting_tool", ImmutableMap.of())); + TestLlm testLlm = createTestLlm(createLlmResponse(functionCall)); + AtomicInteger toolCalls = new AtomicInteger(); + CancellationTokenSource cancellation = new CancellationTokenSource(); + Callbacks.BeforeToolCallback cancelBeforeTool = + (unusedContext, unusedTool, unusedArgs, unusedToolContext) -> { + cancellation.cancel(); + return Maybe.empty(); + }; + LlmAgent agent = + createTestAgentBuilder(testLlm) + .tools(ImmutableList.of(new CountingTool("counting_tool", toolCalls))) + .beforeToolCallback(cancelBeforeTool) + .build(); + InvocationContext invocationContext = + createInvocationContext( + agent, RunConfig.builder().cancellationToken(cancellation.token()).build()); + + List events = + createBaseLlmFlowWithoutProcessors().run(invocationContext).toList().blockingGet(); + + assertThat(events).hasSize(1); + assertEqualIgnoringFunctionIds(events.get(0).content().get(), functionCall); + assertThat(toolCalls.get()).isEqualTo(0); + } + @Test public void run_singleTextResponse_returnsSingleEvent() { Content content = Content.fromParts(Part.fromText("LLM response")); @@ -663,6 +716,26 @@ public Single> runAsync(Map args, ToolContex } } + private static final class CountingTool extends BaseTool { + private final AtomicInteger calls; + + CountingTool(String name, AtomicInteger calls) { + super(name, "tool description for " + name); + this.calls = calls; + } + + @Override + public Optional declaration() { + return Optional.of(FunctionDeclaration.builder().name(name()).build()); + } + + @Override + public Single> runAsync(Map args, ToolContext toolContext) { + calls.incrementAndGet(); + return Single.just(ImmutableMap.of("result", "called")); + } + } + private static class TestLongRunningTool extends BaseTool { private final Map response; diff --git a/core/src/test/java/com/google/adk/runner/RunnerTest.java b/core/src/test/java/com/google/adk/runner/RunnerTest.java index 3870d3461..ba7254541 100644 --- a/core/src/test/java/com/google/adk/runner/RunnerTest.java +++ b/core/src/test/java/com/google/adk/runner/RunnerTest.java @@ -42,6 +42,7 @@ import com.google.adk.agents.BaseAgent; import com.google.adk.agents.Callbacks; import com.google.adk.agents.Callbacks.AfterModelCallback; +import com.google.adk.agents.CancellationTokenSource; import com.google.adk.agents.InvocationContext; import com.google.adk.agents.LiveRequestQueue; import com.google.adk.agents.LlmAgent; @@ -105,6 +106,7 @@ import java.util.UUID; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; import org.jspecify.annotations.Nullable; @@ -178,6 +180,46 @@ public void tearDown() { Tracing.setTracerForTesting(originalTracer); } + @Test + public void runAsync_cancellationCompletesNormallyAndDisposesModelStream() { + AtomicBoolean subscribed = new AtomicBoolean(); + AtomicBoolean disposed = new AtomicBoolean(); + TestLlm neverCompletingLlm = + new TestLlm( + () -> + Flowable.never() + .doOnSubscribe(unused -> subscribed.set(true)) + .doFinally(() -> disposed.set(true))); + Runner cancellationRunner = + Runner.builder() + .app( + App.builder() + .name("cancellation_test") + .rootAgent(createTestAgent(neverCompletingLlm)) + .build()) + .build(); + Session cancellationSession = + cancellationRunner + .sessionService() + .createSession("cancellation_test", "user") + .blockingGet(); + CancellationTokenSource cancellation = new CancellationTokenSource(); + TestSubscriber subscriber = + cancellationRunner + .runAsync( + "user", + cancellationSession.id(), + createContent("stop"), + RunConfig.builder().cancellationToken(cancellation.token()).build()) + .test(); + assertThat(subscribed.get()).isTrue(); + + cancellation.cancel(); + + subscriber.assertComplete().assertNoErrors().assertNoValues(); + assertThat(disposed.get()).isTrue(); + } + @Test public void eventsCompaction_enabled() { TestLlm testLlm =