diff --git a/contrib/langchain4j/src/test/java/com/google/adk/models/langchain4j/LangChain4jIntegrationTest.java b/contrib/langchain4j/src/test/java/com/google/adk/models/langchain4j/LangChain4jIntegrationTest.java index 5b6d3f3ad..ec19339cf 100644 --- a/contrib/langchain4j/src/test/java/com/google/adk/models/langchain4j/LangChain4jIntegrationTest.java +++ b/contrib/langchain4j/src/test/java/com/google/adk/models/langchain4j/LangChain4jIntegrationTest.java @@ -45,7 +45,7 @@ class LangChain4jIntegrationTest { public static final String CLAUDE_4_6_SONNET = "claude-sonnet-4-6"; - public static final String GEMINI_2_0_FLASH = "gemini-2.0-flash"; + public static final String GEMINI_3_6_FLASH = "gemini-3.6-flash"; public static final String GPT_4_O_MINI = "gpt-4o-mini"; @Test @@ -174,7 +174,7 @@ void testAgentTool() { LlmAgent.builder() .name("weather-agent") .description("Weather agent") - .model(GEMINI_2_0_FLASH) + .model(GEMINI_3_6_FLASH) .instruction( """ Your role is to always answer that the weather is sunny and 20°C. @@ -270,7 +270,7 @@ void testSubAgent() { LlmAgent.builder() .name("coordinator-agent") .description("Coordinator agent") - .model(GEMINI_2_0_FLASH) + .model(GEMINI_3_6_FLASH) .instruction( """ Your role is to coordinate 2 agents: @@ -303,8 +303,8 @@ void testSubAgent() { assertEquals(1, hiEvent1.functionCalls().size()); FunctionCall hiFunctionCall = hiEvent1.functionCalls().get(0); assertTrue(hiFunctionCall.id().isPresent()); - assertEquals(Optional.of("transferToAgent"), hiFunctionCall.name()); - assertEquals(Optional.of(Map.of("agentName", "greeterAgent")), hiFunctionCall.args()); + assertEquals(Optional.of("transfer_to_agent"), hiFunctionCall.name()); + assertEquals(Optional.of(Map.of("agent_name", "greeterAgent")), hiFunctionCall.args()); Event hiEvent2 = hiEvents.get(1); assertTrue(hiEvent2.content().isPresent()); @@ -312,7 +312,7 @@ void testSubAgent() { assertEquals(1, hiEvent2.functionResponses().size()); FunctionResponse hiFunctionResponse = hiEvent2.functionResponses().get(0); assertTrue(hiFunctionResponse.id().isPresent()); - assertEquals(Optional.of("transferToAgent"), hiFunctionResponse.name()); + assertEquals(Optional.of("transfer_to_agent"), hiFunctionResponse.name()); assertEquals(Optional.of(Map.of()), hiFunctionResponse.response()); // Empty map for response Event hiEvent3 = hiEvents.get(2); @@ -329,8 +329,8 @@ void testSubAgent() { assertEquals(1, byeEvent1.functionCalls().size()); FunctionCall byeFunctionCall = byeEvent1.functionCalls().get(0); assertTrue(byeFunctionCall.id().isPresent()); - assertEquals(Optional.of("transferToAgent"), byeFunctionCall.name()); - assertEquals(Optional.of(Map.of("agentName", "farewellAgent")), byeFunctionCall.args()); + assertEquals(Optional.of("transfer_to_agent"), byeFunctionCall.name()); + assertEquals(Optional.of(Map.of("agent_name", "farewellAgent")), byeFunctionCall.args()); Event byeEvent2 = byeEvents.get(1); assertTrue(byeEvent2.content().isPresent()); @@ -338,7 +338,7 @@ void testSubAgent() { assertEquals(1, byeEvent2.functionResponses().size()); FunctionResponse byeFunctionResponse = byeEvent2.functionResponses().get(0); assertTrue(byeFunctionResponse.id().isPresent()); - assertEquals(Optional.of("transferToAgent"), byeFunctionResponse.name()); + assertEquals(Optional.of("transfer_to_agent"), byeFunctionResponse.name()); assertEquals(Optional.of(Map.of()), byeFunctionResponse.response()); // Empty map for response Event byeEvent3 = byeEvents.get(2); diff --git a/contrib/spring-ai/src/main/java/com/google/adk/models/springai/MessageConverter.java b/contrib/spring-ai/src/main/java/com/google/adk/models/springai/MessageConverter.java index 442997be0..597e4aaa4 100644 --- a/contrib/spring-ai/src/main/java/com/google/adk/models/springai/MessageConverter.java +++ b/contrib/spring-ai/src/main/java/com/google/adk/models/springai/MessageConverter.java @@ -22,6 +22,7 @@ import com.google.adk.models.LlmResponse; import com.google.genai.types.Content; import com.google.genai.types.FunctionCall; +import com.google.genai.types.FunctionResponse; import com.google.genai.types.GenerateContentResponseUsageMetadata; import com.google.genai.types.Part; import java.net.URI; @@ -54,15 +55,13 @@ * - * - *

Note: Media attachments and tool responses are currently not supported due to Spring AI 1.1.0 - * API limitations (protected/private constructors). These will be added once Spring AI provides - * public APIs for these features. */ public class MessageConverter { + private static final String THOUGHT_SIGNATURES_METADATA_KEY = "thoughtSignatures"; private static final TypeReference> MAP_TYPE_REFERENCE = new TypeReference<>() {}; @@ -261,10 +260,17 @@ private List handleUserContent(Content content) { if (part.text().isPresent()) { textBuilder.append(part.text().get()); } else if (part.functionResponse().isPresent()) { - // TODO: Spring AI 1.1.0 ToolResponseMessage constructors are protected - // For now, we skip tool responses in user messages - // This will need to be addressed in a future update when Spring AI provides - // a public API for creating ToolResponseMessage + FunctionResponse functionResponse = part.functionResponse().get(); + ToolResponseMessage.ToolResponse toolResponse = + new ToolResponseMessage.ToolResponse( + functionResponse.id().orElse(""), + functionResponse + .name() + .orElseThrow( + () -> new IllegalStateException("Function response name is missing")), + toJson(functionResponse.response().orElse(Map.of()))); + toolResponseMessages.add( + ToolResponseMessage.builder().responses(List.of(toolResponse)).build()); } else if (part.inlineData().isPresent()) { // Handle inline media data (images, audio, video, etc.) com.google.genai.types.Blob blob = part.inlineData().get(); @@ -298,8 +304,10 @@ private List handleUserContent(Content content) { } List messages = new ArrayList<>(); - messages.add(UserMessage.builder().text(textBuilder.toString()).media(mediaList).build()); messages.addAll(toolResponseMessages); + if (toolResponseMessages.isEmpty() || textBuilder.length() > 0 || !mediaList.isEmpty()) { + messages.add(UserMessage.builder().text(textBuilder.toString()).media(mediaList).build()); + } return messages; } @@ -307,12 +315,14 @@ private List handleUserContent(Content content) { private AssistantMessage handleAssistantContent(Content content) { StringBuilder textBuilder = new StringBuilder(); List toolCalls = new ArrayList<>(); + List thoughtSignatures = new ArrayList<>(); for (Part part : content.parts().orElse(List.of())) { if (part.text().isPresent()) { textBuilder.append(part.text().get()); } else if (part.functionCall().isPresent()) { FunctionCall functionCall = part.functionCall().get(); + part.thoughtSignature().ifPresent(thoughtSignatures::add); toolCalls.add( new AssistantMessage.ToolCall( functionCall @@ -330,7 +340,12 @@ private AssistantMessage handleAssistantContent(Content content) { if (toolCalls.isEmpty()) { return new AssistantMessage(text); } else { - return AssistantMessage.builder().content(text).toolCalls(toolCalls).build(); + AssistantMessage.Builder builder = + AssistantMessage.builder().content(text).toolCalls(toolCalls); + if (!thoughtSignatures.isEmpty()) { + builder.properties(Map.of(THOUGHT_SIGNATURES_METADATA_KEY, thoughtSignatures)); + } + return builder.build(); } } @@ -434,6 +449,8 @@ private boolean isTurnCompleteResponse(ChatResponse response) { private Content convertAssistantMessageToContent(AssistantMessage assistantMessage) { List parts = new ArrayList<>(); + List thoughtSignatures = getThoughtSignatures(assistantMessage); + int thoughtSignatureIndex = 0; // Add text content if (assistantMessage.getText() != null && !assistantMessage.getText().isEmpty()) { @@ -451,8 +468,13 @@ private Content convertAssistantMessageToContent(AssistantMessage assistantMessa FunctionCall functionCall = FunctionCall.builder().id(toolCall.id()).name(toolCall.name()).args(args).build(); - // Create Part with the FunctionCall (preserves ID) - parts.add(Part.builder().functionCall(functionCall).build()); + // Preserve Gemini thought signatures alongside their corresponding function calls. + Part.Builder partBuilder = Part.builder().functionCall(functionCall); + if (thoughtSignatureIndex < thoughtSignatures.size()) { + partBuilder.thoughtSignature(thoughtSignatures.get(thoughtSignatureIndex)); + thoughtSignatureIndex++; + } + parts.add(partBuilder.build()); } catch (JsonProcessingException e) { throw MessageConversionException.jsonParsingFailed("tool call arguments", e); } @@ -462,6 +484,14 @@ private Content convertAssistantMessageToContent(AssistantMessage assistantMessa return Content.builder().role("model").parts(parts).build(); } + private List getThoughtSignatures(AssistantMessage assistantMessage) { + Object value = assistantMessage.getMetadata().get(THOUGHT_SIGNATURES_METADATA_KEY); + if (!(value instanceof List values)) { + return List.of(); + } + return values.stream().filter(byte[].class::isInstance).map(byte[].class::cast).toList(); + } + private String toJson(Object object) { try { return objectMapper.writeValueAsString(object); diff --git a/contrib/spring-ai/src/test/java/com/google/adk/models/springai/MessageConverterTest.java b/contrib/spring-ai/src/test/java/com/google/adk/models/springai/MessageConverterTest.java index bb529e104..6c84a4694 100644 --- a/contrib/spring-ai/src/test/java/com/google/adk/models/springai/MessageConverterTest.java +++ b/contrib/spring-ai/src/test/java/com/google/adk/models/springai/MessageConverterTest.java @@ -37,6 +37,7 @@ import org.springframework.ai.chat.messages.AssistantMessage; import org.springframework.ai.chat.messages.Message; import org.springframework.ai.chat.messages.SystemMessage; +import org.springframework.ai.chat.messages.ToolResponseMessage; import org.springframework.ai.chat.messages.UserMessage; import org.springframework.ai.chat.metadata.ChatResponseMetadata; import org.springframework.ai.chat.metadata.DefaultUsage; @@ -119,6 +120,7 @@ void testToLlmPromptWithAssistantMessage() { @Test void testToLlmPromptWithFunctionCall() { + byte[] thoughtSignature = {1, 2, 3}; FunctionCall functionCall = FunctionCall.builder() .name("get_weather") @@ -127,7 +129,8 @@ void testToLlmPromptWithFunctionCall() { .build(); // Create Part with FunctionCall inside using Part.builder - Part functionCallPart = Part.builder().functionCall(functionCall).build(); + Part functionCallPart = + Part.builder().functionCall(functionCall).thoughtSignature(thoughtSignature).build(); Content assistantContent = Content.builder() @@ -151,17 +154,12 @@ void testToLlmPromptWithFunctionCall() { assertThat(toolCall.id()).isEqualTo("call_123"); // ID should be preserved now assertThat(toolCall.name()).isEqualTo("get_weather"); assertThat(toolCall.type()).isEqualTo("function"); + assertThat(assistantMessage.getMetadata().get("thoughtSignatures")) + .isEqualTo(List.of(thoughtSignature)); } @Test void testToLlmPromptWithFunctionResponse() { - // TODO: This test is currently limited due to Spring AI 1.1.0 API constraints - // ToolResponseMessage constructors are protected, so function responses are skipped - // Once Spring AI provides public APIs, this test should be updated to verify: - // 1. ToolResponseMessage is created - // 2. Tool response data is properly converted - // 3. Tool call IDs are preserved - FunctionResponse functionResponse = FunctionResponse.builder() .name("get_weather") @@ -174,29 +172,27 @@ void testToLlmPromptWithFunctionResponse() { .role("user") .parts( Part.fromText("What's the weather?"), - Part.fromFunctionResponse( - functionResponse.name().orElse(""), - functionResponse.response().orElse(Map.of()))) + Part.builder().functionResponse(functionResponse).build()) .build(); LlmRequest request = LlmRequest.builder().contents(List.of(userContent)).build(); Prompt prompt = messageConverter.toLlmPrompt(request); - // Currently only UserMessage is created (function response is skipped) - assertThat(prompt.getInstructions()).hasSize(1); + assertThat(prompt.getInstructions()).hasSize(2); - Message userMessage = prompt.getInstructions().get(0); + Message toolResponseMessage = prompt.getInstructions().get(0); + assertThat(toolResponseMessage).isInstanceOf(ToolResponseMessage.class); + ToolResponseMessage toolResponse = (ToolResponseMessage) toolResponseMessage; + assertThat(toolResponse.getResponses()).hasSize(1); + ToolResponseMessage.ToolResponse response = toolResponse.getResponses().get(0); + assertThat(response.id()).isEqualTo("call_123"); + assertThat(response.name()).isEqualTo("get_weather"); + assertThat(response.responseData()).contains("temperature", "72°F", "condition", "sunny"); + + Message userMessage = prompt.getInstructions().get(1); assertThat(userMessage).isInstanceOf(UserMessage.class); assertThat(((UserMessage) userMessage).getText()).isEqualTo("What's the weather?"); - - // When Spring AI provides public API for ToolResponseMessage, uncomment: - // Message toolResponseMessage = prompt.getInstructions().get(1); - // assertThat(toolResponseMessage).isInstanceOf(ToolResponseMessage.class); - // ToolResponseMessage toolResponse = (ToolResponseMessage) toolResponseMessage; - // assertThat(toolResponse.getResponses()).hasSize(1); - // ToolResponseMessage.ToolResponse response = toolResponse.getResponses().get(0); - // assertThat(response.name()).isEqualTo("get_weather"); } @Test @@ -217,6 +213,7 @@ void testToLlmResponseFromChatResponse() { @Test void testToLlmResponseFromChatResponseWithToolCalls() { + byte[] thoughtSignature = {1, 2, 3}; AssistantMessage.ToolCall toolCall = new AssistantMessage.ToolCall( "call_123", "function", "get_weather", "{\"location\":\"San Francisco\"}"); @@ -224,6 +221,7 @@ void testToLlmResponseFromChatResponseWithToolCalls() { AssistantMessage assistantMessage = AssistantMessage.builder() .content("Let me check the weather.") + .properties(Map.of("thoughtSignatures", List.of(thoughtSignature))) .toolCalls(List.of(toolCall)) .build(); @@ -245,6 +243,7 @@ void testToLlmResponseFromChatResponseWithToolCalls() { assertThat(functionCallPart.functionCall().get().name()).contains("get_weather"); // Verify ID is preserved assertThat(functionCallPart.functionCall().get().id()).contains("call_123"); + assertThat(functionCallPart.thoughtSignature()).contains(thoughtSignature); } @Test 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..8c35e01a0 100644 --- a/core/src/main/java/com/google/adk/agents/BaseAgent.java +++ b/core/src/main/java/com/google/adk/agents/BaseAgent.java @@ -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..94032a15d --- /dev/null +++ b/core/src/main/java/com/google/adk/agents/CancellationToken.java @@ -0,0 +1,43 @@ +/* + * 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 io.reactivex.rxjava3.core.Completable; + +/** A thread-safe, cooperative cancellation signal for an agent invocation. */ +@FunctionalInterface +public interface CancellationToken { + /** Returns whether cancellation has been requested. */ + boolean isCancellationRequested(); + + /** Completes when cancellation is requested, or never for polling-only token implementations. */ + default Completable onCancellation() { + return Completable.never(); + } + + /** 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..9d260eb7b --- /dev/null +++ b/core/src/main/java/com/google/adk/agents/CancellationTokenSource.java @@ -0,0 +1,62 @@ +/* + * 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 io.reactivex.rxjava3.core.Completable; +import io.reactivex.rxjava3.subjects.CompletableSubject; +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 CompletableSubject cancellationSignal = CompletableSubject.create(); + private final CancellationToken token = + new CancellationToken() { + @Override + public boolean isCancellationRequested() { + return cancellationRequested.get(); + } + + @Override + public Completable onCancellation() { + return cancellationSignal; + } + }; + + /** 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() { + if (cancellationRequested.compareAndSet(false, true)) { + cancellationSignal.onComplete(); + } + } + + /** 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..a4cd774db 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 @@ -60,8 +60,10 @@ import java.util.List; import java.util.Optional; import java.util.Set; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; import java.util.function.BiFunction; +import java.util.function.Consumer; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -106,7 +108,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 +190,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 +233,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 +265,9 @@ private Flowable callLlm( .switchIfEmpty( Flowable.defer( () -> { + if (context.isCancellationRequested()) { + return Flowable.empty(); + } LlmAgent agent = (LlmAgent) context.agent(); BaseLlm llm = agent.resolvedModel().model().isPresent() @@ -261,6 +280,11 @@ private Flowable callLlm( finalLlmRequest, context.runConfig().streamingMode() == StreamingMode.SSE) + .takeUntil( + context + .cancellationToken() + .onCancellation() + .andThen(Flowable.just(Boolean.TRUE))) .onErrorResumeNext( exception -> handleOnModelErrorCallback( @@ -277,9 +301,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 +335,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 +466,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 +557,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 +601,8 @@ public Flowable runLive(InvocationContext invocationContext) { Flowable.defer( () -> { LlmRequest llmRequestAfterPreprocess = llmRequestRef.get(); - if (invocationContext.endInvocation()) { + if (invocationContext.endInvocation() + || invocationContext.isCancellationRequested()) { return Flowable.empty(); } @@ -582,6 +613,19 @@ public Flowable runLive(InvocationContext invocationContext) { ? agent.resolvedModel().model().get() : LlmRegistry.getLlm(agent.resolvedModel().modelName().get()); BaseLlmConnection connection = llm.connect(llmRequestAfterPreprocess); + AtomicBoolean connectionClosed = new AtomicBoolean(); + Runnable closeConnection = + () -> { + if (connectionClosed.compareAndSet(false, true)) { + connection.close(); + } + }; + Consumer closeConnectionWithError = + error -> { + if (connectionClosed.compareAndSet(false, true)) { + connection.close(error); + } + }; Completable historySent = llmRequestAfterPreprocess.contents().isEmpty() ? Completable.complete() @@ -636,18 +680,18 @@ public Flowable runLive(InvocationContext invocationContext) { } else if (request.blob().isPresent()) { return connection.sendRealtime(request.blob().get()); } - return Completable.fromAction(connection::close); + return Completable.fromAction(closeConnection::run); })) .subscribeWith( new DisposableCompletableObserver() { @Override public void onComplete() { - connection.close(); + closeConnection.run(); } @Override public void onError(Throwable e) { - connection.close(e); + closeConnectionWithError.accept(e); } }); @@ -664,6 +708,12 @@ public void onError(Throwable e) { Flowable receiveFlow = connection .receive() + .takeUntil( + invocationContext + .cancellationToken() + .onCancellation() + .andThen(Flowable.just(Boolean.TRUE))) + .takeWhile(unused -> !invocationContext.isCancellationRequested()) .flatMap( llmResponse -> { Event baseEventForThisLlmResponse = @@ -677,6 +727,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(); @@ -710,9 +763,12 @@ public void onError(Throwable e) { if (event.actions().transferToAgent().isPresent() || event.actions().endInvocation().orElse(false)) { sendTask.dispose(); - connection.close(); + closeConnection.run(); } - }); + }) + .doOnError(closeConnectionWithError::accept) + .doOnCancel(closeConnection::run) + .doFinally(sendTask::dispose); return Tracing.traceFlowable( callLlmContext, 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 =