From f5c6d50571ab3c4938f9ec4ac565c5e517ef9d35 Mon Sep 17 00:00:00 2001 From: Google Team Member Date: Wed, 2 Sep 2026 11:19:44 -0700 Subject: [PATCH] feat(tools): add tool config base, tool error types, and environment simulation PiperOrigin-RevId: 975232714 --- .../adk/flows/llmflows/AgentTransfer.java | 43 +-- .../com/google/adk/tools/BaseToolConfig.java | 51 +++ .../com/google/adk/tools/ToolErrorType.java | 36 ++ .../adk/tools/ToolExecutionException.java | 70 ++++ .../google/adk/tools/TransferToAgentTool.java | 115 ++++++ .../EnvironmentSimulationConfig.java | 113 ++++++ .../EnvironmentSimulationEngine.java | 180 +++++++++ .../EnvironmentSimulationFactory.java | 53 +++ .../EnvironmentSimulationPlugin.java | 40 ++ .../environmentsimulation/InjectedError.java | 60 +++ .../InjectionConfig.java | 110 ++++++ .../environmentsimulation/MockStrategy.java | 30 ++ .../SimulationUtils.java | 75 ++++ .../StatefulParameter.java | 75 ++++ .../ToolConnectionAnalyzer.java | 119 ++++++ .../ToolConnectionMap.java | 57 +++ .../ToolSimulationConfig.java | 78 ++++ .../ToolSpecMockStrategy.java | 254 ++++++++++++ .../google/adk/tools/BaseToolConfigTest.java | 79 ++++ .../adk/tools/ToolExecutionExceptionTest.java | 64 +++ .../adk/tools/TransferToAgentToolTest.java | 83 ++++ .../EnvironmentSimulationConfigTest.java | 121 ++++++ .../EnvironmentSimulationEngineTest.java | 365 ++++++++++++++++++ 23 files changed, 2235 insertions(+), 36 deletions(-) create mode 100644 core/src/main/java/com/google/adk/tools/BaseToolConfig.java create mode 100644 core/src/main/java/com/google/adk/tools/ToolErrorType.java create mode 100644 core/src/main/java/com/google/adk/tools/ToolExecutionException.java create mode 100644 core/src/main/java/com/google/adk/tools/TransferToAgentTool.java create mode 100644 core/src/main/java/com/google/adk/tools/environmentsimulation/EnvironmentSimulationConfig.java create mode 100644 core/src/main/java/com/google/adk/tools/environmentsimulation/EnvironmentSimulationEngine.java create mode 100644 core/src/main/java/com/google/adk/tools/environmentsimulation/EnvironmentSimulationFactory.java create mode 100644 core/src/main/java/com/google/adk/tools/environmentsimulation/EnvironmentSimulationPlugin.java create mode 100644 core/src/main/java/com/google/adk/tools/environmentsimulation/InjectedError.java create mode 100644 core/src/main/java/com/google/adk/tools/environmentsimulation/InjectionConfig.java create mode 100644 core/src/main/java/com/google/adk/tools/environmentsimulation/MockStrategy.java create mode 100644 core/src/main/java/com/google/adk/tools/environmentsimulation/SimulationUtils.java create mode 100644 core/src/main/java/com/google/adk/tools/environmentsimulation/StatefulParameter.java create mode 100644 core/src/main/java/com/google/adk/tools/environmentsimulation/ToolConnectionAnalyzer.java create mode 100644 core/src/main/java/com/google/adk/tools/environmentsimulation/ToolConnectionMap.java create mode 100644 core/src/main/java/com/google/adk/tools/environmentsimulation/ToolSimulationConfig.java create mode 100644 core/src/main/java/com/google/adk/tools/environmentsimulation/ToolSpecMockStrategy.java create mode 100644 core/src/test/java/com/google/adk/tools/BaseToolConfigTest.java create mode 100644 core/src/test/java/com/google/adk/tools/ToolExecutionExceptionTest.java create mode 100644 core/src/test/java/com/google/adk/tools/TransferToAgentToolTest.java create mode 100644 core/src/test/java/com/google/adk/tools/environmentsimulation/EnvironmentSimulationConfigTest.java create mode 100644 core/src/test/java/com/google/adk/tools/environmentsimulation/EnvironmentSimulationEngineTest.java diff --git a/core/src/main/java/com/google/adk/flows/llmflows/AgentTransfer.java b/core/src/main/java/com/google/adk/flows/llmflows/AgentTransfer.java index 0a0da8761..dcb4ceeb3 100644 --- a/core/src/main/java/com/google/adk/flows/llmflows/AgentTransfer.java +++ b/core/src/main/java/com/google/adk/flows/llmflows/AgentTransfer.java @@ -16,17 +16,16 @@ package com.google.adk.flows.llmflows; +import static com.google.common.collect.ImmutableList.toImmutableList; + import com.google.adk.agents.BaseAgent; import com.google.adk.agents.InvocationContext; import com.google.adk.agents.LlmAgent; -import com.google.adk.events.EventActions; import com.google.adk.models.LlmRequest; -import com.google.adk.tools.Annotations.Schema; -import com.google.adk.tools.FunctionTool; import com.google.adk.tools.ToolContext; +import com.google.adk.tools.TransferToAgentTool; import com.google.common.collect.ImmutableList; import io.reactivex.rxjava3.core.Single; -import java.lang.reflect.Method; import java.util.ArrayList; import java.util.List; @@ -55,23 +54,15 @@ public Single processRequest( .appendInstructions( ImmutableList.of(buildTargetAgentsInstructions(agent, transferTargets))); - FunctionTool agentTransferTool = createTransferToAgentTool(); + // Offering only the reachable agents keeps the model from naming one that does not exist. + TransferToAgentTool agentTransferTool = + TransferToAgentTool.create( + transferTargets.stream().map(BaseAgent::name).collect(toImmutableList())); agentTransferTool.processLlmRequest(builder, ToolContext.builder(context).build()); return Single.just( RequestProcessor.RequestProcessingResult.create(builder.build(), ImmutableList.of())); } - private FunctionTool createTransferToAgentTool() { - Method transferToAgentMethod; - try { - transferToAgentMethod = - AgentTransfer.class.getMethod("transferToAgent", String.class, ToolContext.class); - } catch (NoSuchMethodException e) { - throw new IllegalStateException(e); - } - return FunctionTool.create(transferToAgentMethod); - } - /** Builds a string with the target agent’s name and description. */ private String buildTargetAgentsInfo(BaseAgent targetAgent) { return String.format( @@ -144,24 +135,4 @@ private List getTransferTargets(LlmAgent agent) { return transferTargets; } - - @Schema( - name = "transfer_to_agent", - description = - """ - Transfer the question to another agent. - - This tool hands off control to another agent when it's more suitable to - answer the user's question according to the agent's description. - - Args: - agent_name: the agent name to transfer to. - \ - """) - public static void transferToAgent( - @Schema(name = "agent_name") String agentName, - @Schema(optional = true) ToolContext toolContext) { - EventActions eventActions = toolContext.eventActions(); - toolContext.setActions(eventActions.toBuilder().transferToAgent(agentName).build()); - } } diff --git a/core/src/main/java/com/google/adk/tools/BaseToolConfig.java b/core/src/main/java/com/google/adk/tools/BaseToolConfig.java new file mode 100644 index 000000000..0a11b3a07 --- /dev/null +++ b/core/src/main/java/com/google/adk/tools/BaseToolConfig.java @@ -0,0 +1,51 @@ +/* + * 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 + * + * http://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.tools; + +import com.fasterxml.jackson.annotation.JsonAnySetter; +import com.google.adk.JsonBaseModel; + +/** + * The base class for the config of a single tool. + * + *

A tool that takes arguments in a config document gives them a type by extending this class + * with one property per argument, and filling it from the {@link BaseTool.ToolArgsConfig} of the + * entry that named the tool: + * + * 

{@code
+ * MyToolConfig config =
+ *     JsonBaseModel.getMapper().convertValue(args.getAdditionalProperties(), MyToolConfig.class);
+ * }
+ * + *

An argument the subclass does not declare is refused there rather than dropped, so a + * misspelled key in the document reaches the author as an error instead of as an argument that + * silently went missing. + * + *

Experimental. The shape of this type may change. + */ +public abstract class BaseToolConfig extends JsonBaseModel { + + protected BaseToolConfig() {} + + // A key that matches no property of the subclass arrives here, which is where it can be refused: + // the shared mapper skips unknown keys rather than failing on them. + @JsonAnySetter + private void refuseUndeclaredArg(String name, Object value) { + throw new IllegalArgumentException( + String.format("Unknown arg \"%s\" for tool config %s.", name, getClass().getName())); + } +} diff --git a/core/src/main/java/com/google/adk/tools/ToolErrorType.java b/core/src/main/java/com/google/adk/tools/ToolErrorType.java new file mode 100644 index 000000000..a523407f9 --- /dev/null +++ b/core/src/main/java/com/google/adk/tools/ToolErrorType.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 + * + * http://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.tools; + +/** + * How a tool failed, in the HTTP error vocabulary OpenTelemetry semantic conventions use. + * + *

The constant's own name is the reported form: it is what a caller routes on and what fills the + * {@code error.type} span attribute, so it has to read the same here as it does in the other ADK + * languages. + */ +public enum ToolErrorType { + BAD_REQUEST, + UNAUTHORIZED, + FORBIDDEN, + NOT_FOUND, + REQUEST_TIMEOUT, + INTERNAL_SERVER_ERROR, + BAD_GATEWAY, + SERVICE_UNAVAILABLE, + GATEWAY_TIMEOUT +} diff --git a/core/src/main/java/com/google/adk/tools/ToolExecutionException.java b/core/src/main/java/com/google/adk/tools/ToolExecutionException.java new file mode 100644 index 000000000..c556875d2 --- /dev/null +++ b/core/src/main/java/com/google/adk/tools/ToolExecutionException.java @@ -0,0 +1,70 @@ +/* + * 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 + * + * http://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.tools; + +import java.util.Optional; +import org.jspecify.annotations.Nullable; + +/** + * Indicates that a tool failed while it was running. + * + *

A failure carries two things. The message is a sentence a user and a model both end up + * reading, so it should name the tool and say what went wrong. The error type beside it is the + * machine-readable half: it lets a caller tell a timeout apart from a refusal without parsing the + * prose, and it is what fills the {@code error.type} span attribute. Classifying a failure is + * optional, and one that was not classified reports {@link Optional#empty()}. + * + *

No constructor accepts a null error type. A failure with nothing to say about how it failed + * uses {@link #ToolExecutionException(String)} or {@link #ToolExecutionException(String, + * Throwable)} instead. + * + *

This is unchecked, so a tool body can fail inside an RxJava operator and reach the subscriber + * intact. + */ +public class ToolExecutionException extends RuntimeException { + + private final @Nullable ToolErrorType errorType; + + /** Reports a failure with no classification. */ + public ToolExecutionException(String message) { + super(message); + this.errorType = null; + } + + /** Reports a failure with no classification, brought about by {@code cause}. */ + public ToolExecutionException(String message, Throwable cause) { + super(message, cause); + this.errorType = null; + } + + /** Reports a failure classified as {@code errorType}. */ + public ToolExecutionException(String message, ToolErrorType errorType) { + super(message); + this.errorType = errorType; + } + + /** Reports a failure classified as {@code errorType}, brought about by {@code cause}. */ + public ToolExecutionException(String message, ToolErrorType errorType, Throwable cause) { + super(message, cause); + this.errorType = errorType; + } + + /** How the tool failed, empty if the failure was never classified. */ + public Optional errorType() { + return Optional.ofNullable(errorType); + } +} diff --git a/core/src/main/java/com/google/adk/tools/TransferToAgentTool.java b/core/src/main/java/com/google/adk/tools/TransferToAgentTool.java new file mode 100644 index 000000000..802379a84 --- /dev/null +++ b/core/src/main/java/com/google/adk/tools/TransferToAgentTool.java @@ -0,0 +1,115 @@ +/* + * Copyright 2025 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 + * + * http://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.tools; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.genai.types.FunctionDeclaration; +import com.google.genai.types.Schema; +import java.lang.reflect.Method; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; + +/** + * A {@link FunctionTool} that hands the conversation to another agent, restricted to the agents + * that exist. + * + *

The {@code agent_name} parameter carries those agents as an enum in its schema, so a model + * cannot transfer to an agent it was never offered. + * + *

This is the one definition of the {@code transfer_to_agent} tool. The auto flow's transfer + * request processor builds one of these per turn, from the agents that turn can reach, so what the + * model is told about transferring is written in exactly one place. + */ +public final class TransferToAgentTool extends FunctionTool { + + private static final String AGENT_NAME_PARAMETER = "agent_name"; + + private static final Method TRANSFER_TO_AGENT = transferToAgentMethod(); + + private final Optional declaration; + + /** + * Returns a transfer tool that offers the model exactly {@code agentNames} to transfer to. + * + * @param agentNames the valid agent names that can be transferred to. + */ + public static TransferToAgentTool create(List agentNames) { + return new TransferToAgentTool(ImmutableList.copyOf(agentNames)); + } + + private TransferToAgentTool(ImmutableList agentNames) { + super(/* instance= */ null, TRANSFER_TO_AGENT, /* isLongRunning= */ false); + this.declaration = super.declaration().map(decl -> restrictAgentName(decl, agentNames)); + } + + @Override + public Optional declaration() { + return declaration; + } + + /** + * Returns {@code declaration} with its agent name parameter constrained to {@code agentNames}. + */ + private static FunctionDeclaration restrictAgentName( + FunctionDeclaration declaration, ImmutableList agentNames) { + Schema parameters = declaration.parameters().orElse(null); + if (parameters == null) { + return declaration; + } + Map properties = parameters.properties().orElse(ImmutableMap.of()); + Schema agentName = properties.get(AGENT_NAME_PARAMETER); + if (agentName == null) { + return declaration; + } + Map restricted = new LinkedHashMap<>(properties); + restricted.put(AGENT_NAME_PARAMETER, agentName.toBuilder().enum_(agentNames).build()); + return declaration.toBuilder() + .parameters(parameters.toBuilder().properties(restricted).build()) + .build(); + } + + private static Method transferToAgentMethod() { + try { + return TransferToAgentTool.class.getMethod( + "transferToAgent", String.class, ToolContext.class); + } catch (NoSuchMethodException e) { + throw new IllegalStateException(e); + } + } + + @Annotations.Schema( + name = "transfer_to_agent", + description = + """ + Transfer the query to another agent. + + Use this tool to hand off control to another agent that is more suitable to + answer the user's query according to the agent's description. + + Args: + agent_name: the agent name to transfer to. + \ + """) + public static void transferToAgent( + @Annotations.Schema(name = AGENT_NAME_PARAMETER) String agentName, + @Annotations.Schema(optional = true) ToolContext toolContext) { + toolContext.setActions(toolContext.actions().toBuilder().transferToAgent(agentName).build()); + } +} diff --git a/core/src/main/java/com/google/adk/tools/environmentsimulation/EnvironmentSimulationConfig.java b/core/src/main/java/com/google/adk/tools/environmentsimulation/EnvironmentSimulationConfig.java new file mode 100644 index 000000000..bfd704a18 --- /dev/null +++ b/core/src/main/java/com/google/adk/tools/environmentsimulation/EnvironmentSimulationConfig.java @@ -0,0 +1,113 @@ +/* + * 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 + * + * http://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.tools.environmentsimulation; + +import com.google.auto.value.AutoValue; +import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableList; +import com.google.errorprone.annotations.CanIgnoreReturnValue; +import com.google.genai.types.GenerateContentConfig; +import com.google.genai.types.ThinkingConfig; +import java.util.HashSet; +import java.util.Optional; +import java.util.Set; + +/** + * Configuration for an environment simulation: which tools stop really running, and what answers + * them instead. + * + *

Experimental. The shape of this type may change. + */ +@AutoValue +public abstract class EnvironmentSimulationConfig { + + private static final String DEFAULT_SIMULATION_MODEL = "gemini-2.5-flash"; + + /** The tools to simulate. Each tool may be named once. */ + public abstract ImmutableList toolSimulationConfigs(); + + /** + * The model the simulation itself asks, for tool analysis and for mock responses. Defaults to + * {@code gemini-2.5-flash}. + */ + public abstract String simulationModel(); + + /** The configuration for the calls the simulation makes to {@link #simulationModel()}. */ + public abstract GenerateContentConfig simulationModelConfiguration(); + + /** + * Tracing data, such as a prior agent run trace as a JSON string, giving the mock strategy + * historical context. + */ + public abstract Optional tracing(); + + /** + * Environment-specific data, such as a minimal database dump as a JSON string, that the mock + * strategy generates against. + */ + public abstract Optional environmentData(); + + public static Builder builder() { + return new AutoValue_EnvironmentSimulationConfig.Builder() + .simulationModel(DEFAULT_SIMULATION_MODEL) + .simulationModelConfiguration( + GenerateContentConfig.builder() + .thinkingConfig( + ThinkingConfig.builder().includeThoughts(false).thinkingBudget(10240).build()) + .build()); + } + + public abstract Builder toBuilder(); + + /** Builder for {@link EnvironmentSimulationConfig}. */ + @AutoValue.Builder + public abstract static class Builder { + + @CanIgnoreReturnValue + public abstract Builder toolSimulationConfigs( + Iterable toolSimulationConfigs); + + @CanIgnoreReturnValue + public abstract Builder simulationModel(String simulationModel); + + @CanIgnoreReturnValue + public abstract Builder simulationModelConfiguration( + GenerateContentConfig simulationModelConfiguration); + + @CanIgnoreReturnValue + public abstract Builder tracing(String tracing); + + @CanIgnoreReturnValue + public abstract Builder environmentData(String environmentData); + + abstract EnvironmentSimulationConfig autoBuild(); + + public final EnvironmentSimulationConfig build() { + EnvironmentSimulationConfig config = autoBuild(); + Preconditions.checkState( + !config.toolSimulationConfigs().isEmpty(), "toolSimulationConfigs must be provided."); + Set seenToolNames = new HashSet<>(); + for (ToolSimulationConfig toolSimulationConfig : config.toolSimulationConfigs()) { + Preconditions.checkState( + seenToolNames.add(toolSimulationConfig.toolName()), + "Duplicate toolName found: %s. Only one of the entries could ever be reached.", + toolSimulationConfig.toolName()); + } + return config; + } + } +} diff --git a/core/src/main/java/com/google/adk/tools/environmentsimulation/EnvironmentSimulationEngine.java b/core/src/main/java/com/google/adk/tools/environmentsimulation/EnvironmentSimulationEngine.java new file mode 100644 index 000000000..1b48901b0 --- /dev/null +++ b/core/src/main/java/com/google/adk/tools/environmentsimulation/EnvironmentSimulationEngine.java @@ -0,0 +1,180 @@ +/* + * 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 + * + * http://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.tools.environmentsimulation; + +import static com.google.common.collect.ImmutableMap.toImmutableMap; +import static java.util.concurrent.TimeUnit.MILLISECONDS; +import static java.util.function.Function.identity; + +import com.google.adk.agents.LlmAgent; +import com.google.adk.tools.BaseTool; +import com.google.adk.tools.ToolContext; +import com.google.common.collect.ImmutableMap; +import io.reactivex.rxjava3.core.Completable; +import io.reactivex.rxjava3.core.Maybe; +import java.util.Map; +import java.util.Random; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicReference; +import org.jspecify.annotations.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** Decides what answers a tool call that the configuration says should not really run. */ +final class EnvironmentSimulationEngine { + + private static final Logger logger = LoggerFactory.getLogger(EnvironmentSimulationEngine.class); + + private final EnvironmentSimulationConfig config; + private final ImmutableMap toolSimulationConfigs; + private final boolean mocksAnything; + private final ToolConnectionAnalyzer analyzer; + private final ToolSpecMockStrategy mockStrategy; + private final Map>> stateStore = + new ConcurrentHashMap<>(); + private final Random random = new Random(); + + /** The one connection analysis, built on the first tool call that needs it. */ + private final AtomicReference analysis = new AtomicReference<>(); + + private volatile @Nullable ToolConnectionMap toolConnectionMap; + + EnvironmentSimulationEngine(EnvironmentSimulationConfig config) { + this.config = config; + this.toolSimulationConfigs = + config.toolSimulationConfigs().stream() + .collect(toImmutableMap(ToolSimulationConfig::toolName, identity())); + this.mocksAnything = + config.toolSimulationConfigs().stream() + .anyMatch( + toolSimulationConfig -> + toolSimulationConfig.mockStrategyType() + != MockStrategy.MOCK_STRATEGY_UNSPECIFIED); + this.analyzer = + new ToolConnectionAnalyzer(config.simulationModel(), config.simulationModelConfiguration()); + this.mockStrategy = + new ToolSpecMockStrategy(config.simulationModel(), config.simulationModelConfiguration()); + } + + /** + * Simulates one tool call. + * + * @return the response to answer the call with, or empty to let the tool really run. + */ + Maybe> simulate( + BaseTool tool, Map args, ToolContext toolContext) { + ToolSimulationConfig toolSimulationConfig = toolSimulationConfigs.get(tool.name()); + if (toolSimulationConfig == null) { + return Maybe.empty(); + } + return analyzeToolConnections(toolContext) + .andThen(Maybe.defer(() -> inject(toolSimulationConfig, args))) + .switchIfEmpty(Maybe.defer(() -> mock(tool, toolSimulationConfig, args))); + } + + /** + * Waits for the connection analysis, running it if this is the call that has to. + * + *

The analysis is one model call for the whole simulation, but tool calls arrive concurrently, + * so the work is held as a shared cold {@link Completable}: whichever call gets there first + * publishes it, the rest subscribe to that same one, and only one of them ever reaches the model. + * Reading a flag here and setting it when the analysis finished would instead let every call that + * arrived while the first was still running start its own, and the last to finish would decide + * the connection map. + * + *

A configuration that mocks nothing has no use for a connection map, and skips the analysis. + */ + private Completable analyzeToolConnections(ToolContext toolContext) { + if (!mocksAnything) { + return Completable.complete(); + } + return Completable.defer( + () -> + analysis.updateAndGet( + started -> started == null ? analyze(toolContext).cache() : started)); + } + + private Completable analyze(ToolContext toolContext) { + if (!(toolContext.invocationContext().agent() instanceof LlmAgent agent)) { + return Completable.complete(); + } + return agent + .canonicalTools(toolContext) + .toList() + .flatMap(analyzer::analyze) + .doOnSuccess(connectionMap -> toolConnectionMap = connectionMap) + .ignoreElement() + // A simulation is more useful with unconnected mock responses than with none at all, and + // the result is cached, so a failure here must not fail every later tool call too. + .doOnError( + throwable -> + logger.warn( + "Tool connection analysis failed. Proceeding without a connection map.", + throwable)) + .onErrorComplete(); + } + + /** Answers with the first injection whose arguments match and whose probability comes up. */ + private Maybe> inject( + ToolSimulationConfig toolSimulationConfig, Map args) { + for (InjectionConfig injection : toolSimulationConfig.injectionConfigs()) { + if (!args.entrySet().containsAll(injection.matchArgs().entrySet())) { + continue; + } + if (!injects(injection)) { + continue; + } + Maybe> response = + Maybe.just( + injection + .injectedError() + .>map( + error -> + ImmutableMap.of( + "error_code", + error.injectedHttpErrorCode(), + "error_message", + error.errorMessage())) + .orElseGet(injection::injectedResponse)); + double latencySeconds = injection.injectedLatencySeconds(); + return latencySeconds > 0 + ? response.delay((long) (latencySeconds * 1000), MILLISECONDS) + : response; + } + return Maybe.empty(); + } + + /** Seeding and drawing have to be one step, or a concurrent call draws off somebody's seed. */ + private synchronized boolean injects(InjectionConfig injection) { + injection.randomSeed().ifPresent(random::setSeed); + return random.nextDouble() < injection.injectionProbability(); + } + + private Maybe> mock( + BaseTool tool, ToolSimulationConfig toolSimulationConfig, Map args) { + if (toolSimulationConfig.mockStrategyType() == MockStrategy.MOCK_STRATEGY_UNSPECIFIED) { + logger.warn( + "Tool '{}' did not hit any injection config and has no mock strategy configured." + + " Returning no-op.", + tool.name()); + return Maybe.empty(); + } + return mockStrategy + .mock(tool, args, toolConnectionMap, stateStore, config.environmentData(), config.tracing()) + .toMaybe(); + } +} diff --git a/core/src/main/java/com/google/adk/tools/environmentsimulation/EnvironmentSimulationFactory.java b/core/src/main/java/com/google/adk/tools/environmentsimulation/EnvironmentSimulationFactory.java new file mode 100644 index 000000000..d573b4458 --- /dev/null +++ b/core/src/main/java/com/google/adk/tools/environmentsimulation/EnvironmentSimulationFactory.java @@ -0,0 +1,53 @@ +/* + * 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 + * + * http://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.tools.environmentsimulation; + +import com.google.adk.agents.Callbacks.BeforeToolCallback; +import com.google.adk.plugins.Plugin; + +/** + * Creates an environment simulation, either for one agent or for everything a runner runs. + * + *

Experimental. The shape of this type may change. + */ +public final class EnvironmentSimulationFactory { + + /** + * Creates a before-tool callback that simulates the tools the configuration names, for the agent + * it is attached to. + * + * @param config the configuration for the simulation. + * @return a callback to pass to {@code LlmAgent.Builder.beforeToolCallback}. + */ + public static BeforeToolCallback createCallback(EnvironmentSimulationConfig config) { + EnvironmentSimulationEngine engine = new EnvironmentSimulationEngine(config); + return (invocationContext, tool, args, toolContext) -> engine.simulate(tool, args, toolContext); + } + + /** + * Creates a plugin that simulates the tools the configuration names, for every agent the runner + * runs. + * + * @param config the configuration for the simulation. + * @return a plugin to pass to a runner. + */ + public static Plugin createPlugin(EnvironmentSimulationConfig config) { + return new EnvironmentSimulationPlugin(new EnvironmentSimulationEngine(config)); + } + + private EnvironmentSimulationFactory() {} +} diff --git a/core/src/main/java/com/google/adk/tools/environmentsimulation/EnvironmentSimulationPlugin.java b/core/src/main/java/com/google/adk/tools/environmentsimulation/EnvironmentSimulationPlugin.java new file mode 100644 index 000000000..eaf5c015b --- /dev/null +++ b/core/src/main/java/com/google/adk/tools/environmentsimulation/EnvironmentSimulationPlugin.java @@ -0,0 +1,40 @@ +/* + * 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 + * + * http://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.tools.environmentsimulation; + +import com.google.adk.plugins.BasePlugin; +import com.google.adk.tools.BaseTool; +import com.google.adk.tools.ToolContext; +import io.reactivex.rxjava3.core.Maybe; +import java.util.Map; + +/** Simulates the configured tools of every agent the runner runs. */ +final class EnvironmentSimulationPlugin extends BasePlugin { + + private final EnvironmentSimulationEngine engine; + + EnvironmentSimulationPlugin(EnvironmentSimulationEngine engine) { + super("EnvironmentSimulation"); + this.engine = engine; + } + + @Override + public Maybe> beforeToolCallback( + BaseTool tool, Map toolArgs, ToolContext toolContext) { + return engine.simulate(tool, toolArgs, toolContext); + } +} diff --git a/core/src/main/java/com/google/adk/tools/environmentsimulation/InjectedError.java b/core/src/main/java/com/google/adk/tools/environmentsimulation/InjectedError.java new file mode 100644 index 000000000..f408c8fc7 --- /dev/null +++ b/core/src/main/java/com/google/adk/tools/environmentsimulation/InjectedError.java @@ -0,0 +1,60 @@ +/* + * 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 + * + * http://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.tools.environmentsimulation; + +import com.google.auto.value.AutoValue; +import com.google.errorprone.annotations.CanIgnoreReturnValue; + +/** + * An error to be injected into a tool call. + * + *

Experimental. The shape of this type may change. + */ +@AutoValue +public abstract class InjectedError { + + /** + * The HTTP error code to inject into the tool call. Reaches the model as {@code error_code} in + * the tool response. + */ + public abstract int injectedHttpErrorCode(); + + /** + * The error message to inject into the tool call. Reaches the model as {@code error_message} in + * the tool response. + */ + public abstract String errorMessage(); + + public static Builder builder() { + return new AutoValue_InjectedError.Builder(); + } + + public abstract Builder toBuilder(); + + /** Builder for {@link InjectedError}. */ + @AutoValue.Builder + public abstract static class Builder { + + @CanIgnoreReturnValue + public abstract Builder injectedHttpErrorCode(int injectedHttpErrorCode); + + @CanIgnoreReturnValue + public abstract Builder errorMessage(String errorMessage); + + public abstract InjectedError build(); + } +} diff --git a/core/src/main/java/com/google/adk/tools/environmentsimulation/InjectionConfig.java b/core/src/main/java/com/google/adk/tools/environmentsimulation/InjectionConfig.java new file mode 100644 index 000000000..f163cb9ab --- /dev/null +++ b/core/src/main/java/com/google/adk/tools/environmentsimulation/InjectionConfig.java @@ -0,0 +1,110 @@ +/* + * 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 + * + * http://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.tools.environmentsimulation; + +import com.google.auto.value.AutoValue; +import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableMap; +import com.google.errorprone.annotations.CanIgnoreReturnValue; +import java.util.Map; +import java.util.Optional; + +/** + * Injection configuration for a tool: what to answer instead of running it, and when. + * + *

Experimental. The shape of this type may change. + */ +@AutoValue +public abstract class InjectionConfig { + + /** The largest latency that may be injected into a single tool call, in seconds. */ + private static final double MAX_INJECTED_LATENCY_SECONDS = 120.0; + + /** Probability of injecting the injected value. Defaults to 1.0, which always injects. */ + public abstract double injectionProbability(); + + /** + * Injects only into calls whose arguments contain every entry named here. An empty map injects + * into every call. + */ + public abstract ImmutableMap matchArgs(); + + /** + * Latency to inject into the tool call, defaulting to none and capped at 120 seconds. Note it may + * not be accurate if the interceptor is applied as an after-tool callback. + */ + public abstract double injectedLatencySeconds(); + + /** The random seed to use for this injection, which makes the outcome reproducible. */ + public abstract Optional randomSeed(); + + /** The error to answer with. Exactly one of this and {@link #injectedResponse()} is set. */ + public abstract Optional injectedError(); + + /** The response to answer with. Exactly one of this and {@link #injectedError()} is set. */ + public abstract ImmutableMap injectedResponse(); + + public static Builder builder() { + return new AutoValue_InjectionConfig.Builder() + .injectionProbability(1.0) + .matchArgs(ImmutableMap.of()) + .injectedLatencySeconds(0.0) + .injectedResponse(ImmutableMap.of()); + } + + public abstract Builder toBuilder(); + + /** Builder for {@link InjectionConfig}. */ + @AutoValue.Builder + public abstract static class Builder { + + @CanIgnoreReturnValue + public abstract Builder injectionProbability(double injectionProbability); + + @CanIgnoreReturnValue + public abstract Builder matchArgs(Map matchArgs); + + @CanIgnoreReturnValue + public abstract Builder injectedLatencySeconds(double injectedLatencySeconds); + + @CanIgnoreReturnValue + public abstract Builder randomSeed(long randomSeed); + + @CanIgnoreReturnValue + public abstract Builder injectedError(InjectedError injectedError); + + @CanIgnoreReturnValue + public abstract Builder injectedResponse(Map injectedResponse); + + abstract InjectionConfig autoBuild(); + + public final InjectionConfig build() { + InjectionConfig config = autoBuild(); + Preconditions.checkState( + config.injectedLatencySeconds() <= MAX_INJECTED_LATENCY_SECONDS, + "injectedLatencySeconds must be at most %s seconds, but was %s.", + MAX_INJECTED_LATENCY_SECONDS, + config.injectedLatencySeconds()); + boolean hasError = config.injectedError().isPresent(); + boolean hasResponse = !config.injectedResponse().isEmpty(); + Preconditions.checkState( + hasError != hasResponse, + "Either injectedError or injectedResponse must be set, but not both, and not neither."); + return config; + } + } +} diff --git a/core/src/main/java/com/google/adk/tools/environmentsimulation/MockStrategy.java b/core/src/main/java/com/google/adk/tools/environmentsimulation/MockStrategy.java new file mode 100644 index 000000000..97acd71cd --- /dev/null +++ b/core/src/main/java/com/google/adk/tools/environmentsimulation/MockStrategy.java @@ -0,0 +1,30 @@ +/* + * 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 + * + * http://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.tools.environmentsimulation; + +/** + * How a tool call is answered when no injection applies to it. + * + *

Experimental. The shape of this type may change. + */ +public enum MockStrategy { + /** No mock strategy. A tool with no injections either is a configuration mistake. */ + MOCK_STRATEGY_UNSPECIFIED, + + /** Asks the simulation model to write a response from the tool's own schema. */ + MOCK_STRATEGY_TOOL_SPEC +} diff --git a/core/src/main/java/com/google/adk/tools/environmentsimulation/SimulationUtils.java b/core/src/main/java/com/google/adk/tools/environmentsimulation/SimulationUtils.java new file mode 100644 index 000000000..f8cbde146 --- /dev/null +++ b/core/src/main/java/com/google/adk/tools/environmentsimulation/SimulationUtils.java @@ -0,0 +1,75 @@ +/* + * 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 + * + * http://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.tools.environmentsimulation; + +import static java.util.stream.Collectors.joining; + +import com.google.adk.models.LlmRegistry; +import com.google.adk.models.LlmRequest; +import com.google.adk.models.LlmResponse; +import com.google.common.collect.ImmutableList; +import com.google.genai.types.Content; +import com.google.genai.types.GenerateContentConfig; +import com.google.genai.types.Part; +import io.reactivex.rxjava3.core.Single; +import java.util.Optional; + +/** The model call and the JSON handling that the simulation's two model-backed passes share. */ +final class SimulationUtils { + + /** + * Asks the named model for a JSON object and returns the text it wrote. + * + *

The model is looked up by name when the returned {@link Single} is subscribed to, so a + * simulation that never reaches the model never builds one. + */ + static Single generateJson( + String modelName, GenerateContentConfig modelConfig, String prompt) { + return Single.defer( + () -> { + LlmRequest request = + LlmRequest.builder() + .model(modelName) + .contents( + ImmutableList.of( + Content.builder() + .role("user") + .parts(ImmutableList.of(Part.fromText(prompt))) + .build())) + .config(modelConfig.toBuilder().responseMimeType("application/json").build()) + .build(); + return LlmRegistry.getLlm(modelName) + .generateContent(request, /* stream= */ false) + .map(SimulationUtils::text) + .reduce("", String::concat); + }); + } + + /** Strips the code fence a model puts around JSON when it is feeling helpful. */ + static String stripCodeFences(String text) { + return text.replaceAll("^```[a-zA-Z]*\n", "").replaceAll("\n```$", "").trim(); + } + + private static String text(LlmResponse response) { + return response.content().flatMap(Content::parts).orElse(ImmutableList.of()).stream() + .map(Part::text) + .flatMap(Optional::stream) + .collect(joining()); + } + + private SimulationUtils() {} +} diff --git a/core/src/main/java/com/google/adk/tools/environmentsimulation/StatefulParameter.java b/core/src/main/java/com/google/adk/tools/environmentsimulation/StatefulParameter.java new file mode 100644 index 000000000..e2e2b4cde --- /dev/null +++ b/core/src/main/java/com/google/adk/tools/environmentsimulation/StatefulParameter.java @@ -0,0 +1,75 @@ +/* + * 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 + * + * http://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.tools.environmentsimulation; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.fasterxml.jackson.databind.annotation.JsonDeserialize; +import com.fasterxml.jackson.databind.annotation.JsonPOJOBuilder; +import com.google.adk.JsonBaseModel; +import com.google.auto.value.AutoValue; +import com.google.common.collect.ImmutableList; +import com.google.errorprone.annotations.CanIgnoreReturnValue; + +/** A parameter shared between tools, and the tools on each side of it. */ +@AutoValue +@JsonDeserialize(builder = StatefulParameter.Builder.class) +abstract class StatefulParameter extends JsonBaseModel { + + /** The name of the shared parameter, for example {@code ticket_id}. */ + @JsonProperty("parameter_name") + abstract String parameterName(); + + /** The tools that generate this parameter. */ + @JsonProperty("creating_tools") + abstract ImmutableList creatingTools(); + + /** The tools that take this parameter as input. */ + @JsonProperty("consuming_tools") + abstract ImmutableList consumingTools(); + + static Builder builder() { + return new AutoValue_StatefulParameter.Builder() + .creatingTools(ImmutableList.of()) + .consumingTools(ImmutableList.of()); + } + + /** Builder for {@link StatefulParameter}. */ + @AutoValue.Builder + @JsonPOJOBuilder(buildMethodName = "build", withPrefix = "") + abstract static class Builder { + + @JsonCreator + static Builder jacksonBuilder() { + return StatefulParameter.builder(); + } + + @CanIgnoreReturnValue + @JsonProperty("parameter_name") + abstract Builder parameterName(String parameterName); + + @CanIgnoreReturnValue + @JsonProperty("creating_tools") + abstract Builder creatingTools(Iterable creatingTools); + + @CanIgnoreReturnValue + @JsonProperty("consuming_tools") + abstract Builder consumingTools(Iterable consumingTools); + + abstract StatefulParameter build(); + } +} diff --git a/core/src/main/java/com/google/adk/tools/environmentsimulation/ToolConnectionAnalyzer.java b/core/src/main/java/com/google/adk/tools/environmentsimulation/ToolConnectionAnalyzer.java new file mode 100644 index 000000000..50a31bca7 --- /dev/null +++ b/core/src/main/java/com/google/adk/tools/environmentsimulation/ToolConnectionAnalyzer.java @@ -0,0 +1,119 @@ +/* + * 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 + * + * http://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.tools.environmentsimulation; + +import static java.util.stream.Collectors.joining; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.google.adk.JsonBaseModel; +import com.google.adk.tools.BaseTool; +import com.google.genai.types.FunctionDeclaration; +import com.google.genai.types.GenerateContentConfig; +import io.reactivex.rxjava3.core.Single; +import java.util.List; +import java.util.Optional; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Asks a model which parameters the tools share, so that a mock response for one tool can stay + * consistent with the mock responses of the tools around it. For example, {@code get_ticket} + * consumes a {@code ticket_id} that {@code create_ticket} produced. + */ +final class ToolConnectionAnalyzer { + + private static final Logger logger = LoggerFactory.getLogger(ToolConnectionAnalyzer.class); + + private static final String PROMPT_TEMPLATE = + """ + You are an expert software architect analyzing a set of tools to understand + stateful dependencies. Your task is to identify parameters that act as + stateful identifiers (like IDs) and classify the tools that interact with + them. + + **Definitions:** + - A **"creating tool"** is a tool that creates a new resource or makes a + significant state change to an existing one (e.g., creating, updating, + canceling, or deleting). Tool names like `create_account`, `cancel_order`, + or `update_price` are strong indicators. These tools are responsible for + generating or modifying the state associated with an ID. + - A **"consuming tool"** is a tool that uses a resource's ID to retrieve + information without changing its state. Tool names like `get_user`, + `list_events`, or `find_order` are strong indicators. + + **Your Goal:** + Analyze the following tool schemas and identify the shared, stateful + parameters (like `user_id`, `order_id`, etc.). + + For each stateful parameter you identify, classify the tools into + `creating_tools` and `consuming_tools` based on the definitions above. + + **Example:** A `create_ticket` tool would be a `creating_tool` for + `ticket_id`. A `get_ticket` tool would be a `consuming_tool` for + `ticket_id`. A `list_tickets` tool that takes a `user_id` as input is a + `consuming_tool` for `user_id`. + + **Analyze the following tool schemas:** + {tool_schemas_json} + + **Output Format:** + Generate a JSON object with a single key, "stateful_parameters", which is a + list. Each item in the list must have these keys: + - "parameter_name": The name of the shared parameter (e.g., "ticket_id"). + - "creating_tools": A list of tools that create or modify this parameter's + state. + - "consuming_tools": A list of tools that use this parameter as input for + read-only operations. + + ONLY return the raw JSON object. + Your response must start with '{' and end with '}'. + """; + + private final String modelName; + private final GenerateContentConfig modelConfig; + + ToolConnectionAnalyzer(String modelName, GenerateContentConfig modelConfig) { + this.modelName = modelName; + this.modelConfig = modelConfig; + } + + /** Analyzes the given tools and returns the map of their connections. */ + Single analyze(List tools) { + String toolSchemasJson = + tools.stream() + .map(BaseTool::declaration) + .flatMap(Optional::stream) + .map(FunctionDeclaration::toJson) + .collect(joining(",\n", "[\n", "\n]")); + String prompt = PROMPT_TEMPLATE.replace("{tool_schemas_json}", toolSchemasJson); + return SimulationUtils.generateJson(modelName, modelConfig, prompt).map(this::parse); + } + + private ToolConnectionMap parse(String responseText) { + try { + return JsonBaseModel.getMapper() + .readValue(SimulationUtils.stripCodeFences(responseText), ToolConnectionMap.class); + } catch (JsonProcessingException e) { + logger.warn( + "Failed to read a tool connection analysis from the model. Proceeding without a" + + " connection map. Model output:\n{}", + responseText, + e); + return ToolConnectionMap.builder().build(); + } + } +} diff --git a/core/src/main/java/com/google/adk/tools/environmentsimulation/ToolConnectionMap.java b/core/src/main/java/com/google/adk/tools/environmentsimulation/ToolConnectionMap.java new file mode 100644 index 000000000..94a8d8010 --- /dev/null +++ b/core/src/main/java/com/google/adk/tools/environmentsimulation/ToolConnectionMap.java @@ -0,0 +1,57 @@ +/* + * 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 + * + * http://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.tools.environmentsimulation; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.fasterxml.jackson.databind.annotation.JsonDeserialize; +import com.fasterxml.jackson.databind.annotation.JsonPOJOBuilder; +import com.google.adk.JsonBaseModel; +import com.google.auto.value.AutoValue; +import com.google.common.collect.ImmutableList; +import com.google.errorprone.annotations.CanIgnoreReturnValue; + +/** How the simulated tools connect to one another through shared, stateful parameters. */ +@AutoValue +@JsonDeserialize(builder = ToolConnectionMap.Builder.class) +abstract class ToolConnectionMap extends JsonBaseModel { + + /** The stateful parameters and their connections. */ + @JsonProperty("stateful_parameters") + abstract ImmutableList statefulParameters(); + + static Builder builder() { + return new AutoValue_ToolConnectionMap.Builder().statefulParameters(ImmutableList.of()); + } + + /** Builder for {@link ToolConnectionMap}. */ + @AutoValue.Builder + @JsonPOJOBuilder(buildMethodName = "build", withPrefix = "") + abstract static class Builder { + + @JsonCreator + static Builder jacksonBuilder() { + return ToolConnectionMap.builder(); + } + + @CanIgnoreReturnValue + @JsonProperty("stateful_parameters") + abstract Builder statefulParameters(Iterable statefulParameters); + + abstract ToolConnectionMap build(); + } +} diff --git a/core/src/main/java/com/google/adk/tools/environmentsimulation/ToolSimulationConfig.java b/core/src/main/java/com/google/adk/tools/environmentsimulation/ToolSimulationConfig.java new file mode 100644 index 000000000..2b18b57b7 --- /dev/null +++ b/core/src/main/java/com/google/adk/tools/environmentsimulation/ToolSimulationConfig.java @@ -0,0 +1,78 @@ +/* + * 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 + * + * http://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.tools.environmentsimulation; + +import com.google.auto.value.AutoValue; +import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableList; +import com.google.errorprone.annotations.CanIgnoreReturnValue; + +/** + * Simulation configuration for a single tool. + * + *

Experimental. The shape of this type may change. + */ +@AutoValue +public abstract class ToolSimulationConfig { + + /** Name of the tool to be simulated. */ + public abstract String toolName(); + + /** + * Injections for the tool, tried in order. The mock strategy answers the call when none of them + * applies. + */ + public abstract ImmutableList injectionConfigs(); + + /** The mock strategy to use. */ + public abstract MockStrategy mockStrategyType(); + + public static Builder builder() { + return new AutoValue_ToolSimulationConfig.Builder() + .injectionConfigs(ImmutableList.of()) + .mockStrategyType(MockStrategy.MOCK_STRATEGY_UNSPECIFIED); + } + + public abstract Builder toBuilder(); + + /** Builder for {@link ToolSimulationConfig}. */ + @AutoValue.Builder + public abstract static class Builder { + + @CanIgnoreReturnValue + public abstract Builder toolName(String toolName); + + @CanIgnoreReturnValue + public abstract Builder injectionConfigs(Iterable injectionConfigs); + + @CanIgnoreReturnValue + public abstract Builder mockStrategyType(MockStrategy mockStrategyType); + + abstract ToolSimulationConfig autoBuild(); + + public final ToolSimulationConfig build() { + ToolSimulationConfig config = autoBuild(); + Preconditions.checkState( + !config.injectionConfigs().isEmpty() + || config.mockStrategyType() != MockStrategy.MOCK_STRATEGY_UNSPECIFIED, + "Tool \"%s\" has no injectionConfigs, so mockStrategyType cannot be" + + " MOCK_STRATEGY_UNSPECIFIED: nothing would ever be simulated for it.", + config.toolName()); + return config; + } + } +} diff --git a/core/src/main/java/com/google/adk/tools/environmentsimulation/ToolSpecMockStrategy.java b/core/src/main/java/com/google/adk/tools/environmentsimulation/ToolSpecMockStrategy.java new file mode 100644 index 000000000..178356164 --- /dev/null +++ b/core/src/main/java/com/google/adk/tools/environmentsimulation/ToolSpecMockStrategy.java @@ -0,0 +1,254 @@ +/* + * 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 + * + * http://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.tools.environmentsimulation; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.core.type.TypeReference; +import com.google.adk.JsonBaseModel; +import com.google.adk.tools.BaseTool; +import com.google.common.collect.ImmutableMap; +import com.google.genai.types.FunctionDeclaration; +import com.google.genai.types.GenerateContentConfig; +import io.reactivex.rxjava3.core.Single; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.ConcurrentHashMap; +import java.util.regex.Matcher; +import java.util.regex.Pattern; +import org.jspecify.annotations.Nullable; + +/** Answers a tool call with a response a model writes from the tool's own specification. */ +final class ToolSpecMockStrategy { + + private static final Pattern PLACEHOLDER_PATTERN = Pattern.compile("\\{(\\w+)\\}"); + + private static final String PROMPT_TEMPLATE = + """ + You are a stateful tool simulator. Your task is to generate a + realistic JSON response for a tool call, maintaining consistency based + on a shared state. + + {environment_data_snippet} + + {tracing_snippet} + + Here is the map of how tools connect via stateful parameters: + {tool_connection_map_json} + + Here is the current state of all stateful parameters: + {state_store_json} + + You are now simulating the following tool call: + Tool Name: {tool_name} + Tool Description: {tool_description} + Tool Schema: {tool_schema_json} + Tool Arguments: {tool_arguments_json} + + Your instructions: + 1. Analyze the tool call. Is it a "creating" or "consuming" tool + based on the connection map? + 2. If it's a "consuming" tool, check the provided arguments against + the state store. If an ID is provided that does not exist in the + state, return a realistic error (e.g., a 404 Not Found error). + Otherwise, use the data from the state, the provided environment data, + and the tracing history to generate the response. + 3. If it's a "creating" tool, generate a new, unique ID for the + stateful parameter (e.g., a random string for a ticket_id). Include + this new ID in your response. I will then update the state with it. + 4. Leverage the provided environment data (if any) to make your response + more realistic and consistent with the simulated environment. + 5. Leverage the provided tracing history (if any) to make your response + consistent with observed tool behavior patterns from prior runs. + 6. Generate a convincing, valid JSON object that mocks the tool's + response. The response must be only the JSON object, without any + additional text or formatting. + 7. The response must start with '{' and end with '}'. + """; + + private static final String ENVIRONMENT_DATA_SNIPPET_TEMPLATE = + """ + Here is relevant environment data (e.g., database snippet, context information): + + {environment_data} + + Use this information to generate more realistic responses. + """; + + private static final String TRACING_SNIPPET_TEMPLATE = + """ + Here is a tracing history from a prior agent run (e.g., recorded tool + calls and responses): + + {tracing} + + Use this history to make your mock responses consistent with observed + tool behavior patterns. + """; + + private final String modelName; + private final GenerateContentConfig modelConfig; + + ToolSpecMockStrategy(String modelName, GenerateContentConfig modelConfig) { + this.modelName = modelName; + this.modelConfig = modelConfig; + } + + /** + * Generates a mock response for one tool call, and records anything the tool created in the state + * store so that later calls stay consistent with it. + */ + Single> mock( + BaseTool tool, + Map args, + @Nullable ToolConnectionMap toolConnectionMap, + Map>> stateStore, + Optional environmentData, + Optional tracing) { + Optional declaration = tool.declaration(); + if (declaration.isEmpty()) { + return Single.just( + ImmutableMap.of("status", "error", "error_message", "Could not get tool declaration.")); + } + + String prompt = + fill( + PROMPT_TEMPLATE, + ImmutableMap.builder() + .put( + "environment_data_snippet", + environmentData + .map( + data -> + fill( + ENVIRONMENT_DATA_SNIPPET_TEMPLATE, + ImmutableMap.of("environment_data", data))) + .orElse("")) + .put( + "tracing_snippet", + tracing + .map( + trace -> + fill(TRACING_SNIPPET_TEMPLATE, ImmutableMap.of("tracing", trace))) + .orElse("")) + .put( + "tool_connection_map_json", + toolConnectionMap == null ? "''" : toolConnectionMap.toJson()) + .put("state_store_json", JsonBaseModel.toJsonString(stateStore)) + .put("tool_name", tool.name()) + .put("tool_description", tool.description()) + .put("tool_schema_json", declaration.get().toJson()) + .put("tool_arguments_json", JsonBaseModel.toJsonString(args)) + .buildOrThrow()); + + return SimulationUtils.generateJson(modelName, modelConfig, prompt) + .map(responseText -> toMockResponse(tool, responseText, toolConnectionMap, stateStore)); + } + + /** + * Substitutes every {@code {name}} placeholder in one left-to-right pass. A value is written to + * the output and never scanned again, so a placeholder that happens to appear inside environment + * data, a tracing history or a tool's own description stays literal text rather than becoming a + * substitution site. A placeholder with no value keeps its braces. + */ + private static String fill(String template, Map values) { + Matcher matcher = PLACEHOLDER_PATTERN.matcher(template); + StringBuilder filled = new StringBuilder(); + int copiedUpTo = 0; + while (matcher.find()) { + String value = values.get(matcher.group(1)); + if (value == null) { + continue; + } + filled.append(template, copiedUpTo, matcher.start()).append(value); + copiedUpTo = matcher.end(); + } + return filled.append(template, copiedUpTo, template.length()).toString(); + } + + private static Map toMockResponse( + BaseTool tool, + String responseText, + @Nullable ToolConnectionMap toolConnectionMap, + Map>> stateStore) { + Map mockResponse; + try { + mockResponse = + JsonBaseModel.getMapper() + .readValue( + SimulationUtils.stripCodeFences(responseText), + new TypeReference>() {}); + } catch (JsonProcessingException e) { + return ImmutableMap.of( + "status", + "error", + "error_message", + "Failed to generate valid JSON mock response.", + "llm_output", + responseText); + } + recordCreatedState(tool, mockResponse, toolConnectionMap, stateStore); + return mockResponse; + } + + /** + * Stores the response under every stateful parameter this tool creates, which is what lets a + * later call that consumes the same parameter be answered consistently. + */ + private static void recordCreatedState( + BaseTool tool, + Map mockResponse, + @Nullable ToolConnectionMap toolConnectionMap, + Map>> stateStore) { + if (toolConnectionMap == null) { + return; + } + for (StatefulParameter parameter : toolConnectionMap.statefulParameters()) { + if (!parameter.creatingTools().contains(tool.name())) { + continue; + } + Object parameterValue = findValueByKey(mockResponse, parameter.parameterName()); + if (parameterValue != null) { + stateStore + .computeIfAbsent(parameter.parameterName(), name -> new ConcurrentHashMap<>()) + .put(parameterValue, mockResponse); + } + } + } + + private static @Nullable Object findValueByKey(Object data, String targetKey) { + if (data instanceof Map map) { + if (map.containsKey(targetKey)) { + return map.get(targetKey); + } + for (Object value : map.values()) { + Object result = findValueByKey(value, targetKey); + if (result != null) { + return result; + } + } + } else if (data instanceof List list) { + for (Object item : list) { + Object result = findValueByKey(item, targetKey); + if (result != null) { + return result; + } + } + } + return null; + } +} diff --git a/core/src/test/java/com/google/adk/tools/BaseToolConfigTest.java b/core/src/test/java/com/google/adk/tools/BaseToolConfigTest.java new file mode 100644 index 000000000..09dec6567 --- /dev/null +++ b/core/src/test/java/com/google/adk/tools/BaseToolConfigTest.java @@ -0,0 +1,79 @@ +/* + * 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 + * + * http://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.tools; + +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertThrows; + +import com.google.adk.JsonBaseModel; +import com.google.common.collect.ImmutableMap; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public final class BaseToolConfigTest { + + /** Stands in for the config of a tool that takes one argument. */ + public static final class WeatherToolConfig extends BaseToolConfig { + private String city; + + // Bound by name through this setter, which is how the deserializer discovers the property + // without a binding annotation. + public void setCity(String city) { + this.city = city; + } + + public String city() { + return city; + } + } + + @Test + public void declaredArg_isRead() { + WeatherToolConfig config = + JsonBaseModel.getMapper() + .convertValue(ImmutableMap.of("city", "Zurich"), WeatherToolConfig.class); + + assertThat(config.city()).isEqualTo("Zurich"); + } + + @Test + public void undeclaredArg_isRefused() { + // The shared mapper drops unknown keys rather than failing on them, so without the base class + // refusing them this misspelling would silently read back as no city at all. + IllegalArgumentException exception = + assertThrows( + IllegalArgumentException.class, + () -> + JsonBaseModel.getMapper() + .convertValue(ImmutableMap.of("citty", "Zurich"), WeatherToolConfig.class)); + + assertThat(exception).hasMessageThat().contains("citty"); + } + + @Test + public void undeclaredArgAlongsideDeclaredOne_isRefused() { + assertThrows( + IllegalArgumentException.class, + () -> + JsonBaseModel.getMapper() + .convertValue( + ImmutableMap.of("city", "Zurich", "units", "celsius"), + WeatherToolConfig.class)); + } +} diff --git a/core/src/test/java/com/google/adk/tools/ToolExecutionExceptionTest.java b/core/src/test/java/com/google/adk/tools/ToolExecutionExceptionTest.java new file mode 100644 index 000000000..303b885ed --- /dev/null +++ b/core/src/test/java/com/google/adk/tools/ToolExecutionExceptionTest.java @@ -0,0 +1,64 @@ +/* + * 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 + * + * http://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.tools; + +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 ToolExecutionExceptionTest { + + @Test + public void errorType_notClassified_isEmpty() { + ToolExecutionException exception = new ToolExecutionException("weather_tool gave up."); + + assertThat(exception).hasMessageThat().isEqualTo("weather_tool gave up."); + assertThat(exception.errorType()).isEmpty(); + } + + @Test + public void errorType_classified_isWhatWasGiven() { + ToolExecutionException exception = + new ToolExecutionException("weather_tool timed out.", ToolErrorType.GATEWAY_TIMEOUT); + + assertThat(exception.errorType()).hasValue(ToolErrorType.GATEWAY_TIMEOUT); + } + + @Test + public void cause_withoutClassification_isKept() { + Exception cause = new IllegalStateException("socket closed"); + + ToolExecutionException exception = new ToolExecutionException("weather_tool failed.", cause); + + assertThat(exception).hasCauseThat().isSameInstanceAs(cause); + assertThat(exception.errorType()).isEmpty(); + } + + @Test + public void cause_withClassification_isKept() { + Exception cause = new IllegalStateException("socket closed"); + + ToolExecutionException exception = + new ToolExecutionException("weather_tool failed.", ToolErrorType.BAD_GATEWAY, cause); + + assertThat(exception).hasCauseThat().isSameInstanceAs(cause); + assertThat(exception.errorType()).hasValue(ToolErrorType.BAD_GATEWAY); + } +} diff --git a/core/src/test/java/com/google/adk/tools/TransferToAgentToolTest.java b/core/src/test/java/com/google/adk/tools/TransferToAgentToolTest.java new file mode 100644 index 000000000..8e8ea2f51 --- /dev/null +++ b/core/src/test/java/com/google/adk/tools/TransferToAgentToolTest.java @@ -0,0 +1,83 @@ +/* + * 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 + * + * http://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.tools; + +import static com.google.adk.testing.TestUtils.createInvocationContext; +import static com.google.adk.testing.TestUtils.createRootAgent; +import static com.google.common.truth.Truth.assertThat; + +import com.google.adk.agents.InvocationContext; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.genai.types.Schema; +import java.util.Map; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public final class TransferToAgentToolTest { + + @Test + public void name_isTheNameTheFlowAndTheModelAgreeOn() { + TransferToAgentTool tool = TransferToAgentTool.create(ImmutableList.of("billing_agent")); + + assertThat(tool.name()).isEqualTo("transfer_to_agent"); + } + + @Test + public void declaration_offersOnlyTheGivenAgents() { + TransferToAgentTool tool = + TransferToAgentTool.create(ImmutableList.of("billing_agent", "support_agent")); + + assertThat(agentNameSchema(tool).enum_()) + .hasValue(ImmutableList.of("billing_agent", "support_agent")); + } + + @Test + public void declaration_withNoAgents_offersNothing() { + TransferToAgentTool tool = TransferToAgentTool.create(ImmutableList.of()); + + assertThat(agentNameSchema(tool).enum_()).hasValue(ImmutableList.of()); + } + + @Test + public void declaration_leavesTheRestOfTheSchemaAlone() { + TransferToAgentTool tool = TransferToAgentTool.create(ImmutableList.of("billing_agent")); + + assertThat(tool.declaration().get().parameters().get().required()) + .hasValue(ImmutableList.of("agent_name")); + assertThat(tool.declaration().get().description()).isPresent(); + } + + @Test + public void runAsync_recordsTheTransferOnTheContext() { + TransferToAgentTool tool = + TransferToAgentTool.create(ImmutableList.of("billing_agent", "support_agent")); + InvocationContext invocationContext = createInvocationContext(createRootAgent()); + ToolContext toolContext = ToolContext.builder(invocationContext).build(); + + Map unused = + tool.runAsync(ImmutableMap.of("agent_name", "billing_agent"), toolContext).blockingGet(); + + assertThat(toolContext.actions().transferToAgent()).hasValue("billing_agent"); + } + + private static Schema agentNameSchema(TransferToAgentTool tool) { + return tool.declaration().get().parameters().get().properties().get().get("agent_name"); + } +} diff --git a/core/src/test/java/com/google/adk/tools/environmentsimulation/EnvironmentSimulationConfigTest.java b/core/src/test/java/com/google/adk/tools/environmentsimulation/EnvironmentSimulationConfigTest.java new file mode 100644 index 000000000..63c515b98 --- /dev/null +++ b/core/src/test/java/com/google/adk/tools/environmentsimulation/EnvironmentSimulationConfigTest.java @@ -0,0 +1,121 @@ +/* + * 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 + * + * http://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.tools.environmentsimulation; + +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertThrows; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public final class EnvironmentSimulationConfigTest { + + @Test + public void build_duplicateToolName_isRefused() { + EnvironmentSimulationConfig.Builder builder = + EnvironmentSimulationConfig.builder() + .toolSimulationConfigs( + ImmutableList.of(mockingTool("weather"), mockingTool("weather"))); + + IllegalStateException exception = assertThrows(IllegalStateException.class, builder::build); + + assertThat(exception).hasMessageThat().contains("weather"); + } + + @Test + public void build_noToolSimulationConfigs_isRefused() { + EnvironmentSimulationConfig.Builder builder = + EnvironmentSimulationConfig.builder().toolSimulationConfigs(ImmutableList.of()); + + assertThrows(IllegalStateException.class, builder::build); + } + + @Test + public void build_defaultsToAModelAndAThinkingBudget() { + EnvironmentSimulationConfig config = + EnvironmentSimulationConfig.builder() + .toolSimulationConfigs(ImmutableList.of(mockingTool("weather"))) + .build(); + + assertThat(config.simulationModel()).isEqualTo("gemini-2.5-flash"); + assertThat(config.simulationModelConfiguration().thinkingConfig()).isPresent(); + assertThat(config.tracing()).isEmpty(); + assertThat(config.environmentData()).isEmpty(); + } + + @Test + public void toolSimulationConfig_nothingToSimulate_isRefused() { + // No injections and no mock strategy means the tool is named but never simulated, which is a + // configuration mistake rather than a way to switch simulation off for it. + ToolSimulationConfig.Builder builder = ToolSimulationConfig.builder().toolName("weather"); + + IllegalStateException exception = assertThrows(IllegalStateException.class, builder::build); + + assertThat(exception).hasMessageThat().contains("weather"); + } + + @Test + public void injectionConfig_bothAnErrorAndAResponse_isRefused() { + InjectionConfig.Builder builder = + InjectionConfig.builder() + .injectedError( + InjectedError.builder() + .injectedHttpErrorCode(503) + .errorMessage("upstream is down") + .build()) + .injectedResponse(ImmutableMap.of("temperature", 20)); + + assertThrows(IllegalStateException.class, builder::build); + } + + @Test + public void injectionConfig_neitherAnErrorNorAResponse_isRefused() { + assertThrows(IllegalStateException.class, InjectionConfig.builder()::build); + } + + @Test + public void injectionConfig_latencyOverTheCap_isRefused() { + InjectionConfig.Builder builder = + InjectionConfig.builder() + .injectedLatencySeconds(120.1) + .injectedResponse(ImmutableMap.of("temperature", 20)); + + assertThrows(IllegalStateException.class, builder::build); + } + + @Test + public void injectionConfig_defaultsToAlwaysInjectingEveryCall() { + InjectionConfig config = + InjectionConfig.builder().injectedResponse(ImmutableMap.of("temperature", 20)).build(); + + assertThat(config.injectionProbability()).isEqualTo(1.0); + assertThat(config.matchArgs()).isEmpty(); + assertThat(config.injectedLatencySeconds()).isEqualTo(0.0); + assertThat(config.randomSeed()).isEmpty(); + } + + private static ToolSimulationConfig mockingTool(String toolName) { + return ToolSimulationConfig.builder() + .toolName(toolName) + .mockStrategyType(MockStrategy.MOCK_STRATEGY_TOOL_SPEC) + .build(); + } +} diff --git a/core/src/test/java/com/google/adk/tools/environmentsimulation/EnvironmentSimulationEngineTest.java b/core/src/test/java/com/google/adk/tools/environmentsimulation/EnvironmentSimulationEngineTest.java new file mode 100644 index 000000000..a6337c58c --- /dev/null +++ b/core/src/test/java/com/google/adk/tools/environmentsimulation/EnvironmentSimulationEngineTest.java @@ -0,0 +1,365 @@ +/* + * 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 + * + * http://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.tools.environmentsimulation; + +import static com.google.adk.testing.TestUtils.createInvocationContext; +import static com.google.adk.testing.TestUtils.createTestAgentBuilder; +import static com.google.adk.testing.TestUtils.createTestLlm; +import static com.google.adk.testing.TestUtils.createTextLlmResponse; +import static com.google.common.truth.Truth.assertThat; +import static java.util.concurrent.TimeUnit.MILLISECONDS; +import static java.util.concurrent.TimeUnit.SECONDS; +import static java.util.stream.Collectors.joining; + +import com.google.adk.agents.LlmAgent; +import com.google.adk.models.LlmRegistry; +import com.google.adk.models.LlmResponse; +import com.google.adk.testing.TestLlm; +import com.google.adk.tools.BaseTool; +import com.google.adk.tools.ToolContext; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.genai.types.FunctionDeclaration; +import io.reactivex.rxjava3.core.Flowable; +import java.time.Duration; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Random; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; +import java.util.concurrent.CyclicBarrier; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import org.junit.Before; +import org.junit.BeforeClass; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TestName; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public final class EnvironmentSimulationEngineTest { + + /** A phrase that appears only in the prompt the connection analysis sends. */ + private static final String ANALYSIS_PROMPT = "expert software architect"; + + /** A phrase that appears only in the prompt a mocked tool response is generated from. */ + private static final String MOCK_PROMPT = "stateful tool simulator"; + + private static final String MOCK_RESPONSE_JSON = "{\"status\": \"ok\"}"; + + // The registry caches one model instance per name, so each test gets its own name and finds its + // own fake here rather than the one a test that ran earlier left behind. + private static final ConcurrentMap SIMULATION_MODELS = new ConcurrentHashMap<>(); + + @BeforeClass + public static void registerSimulationModels() { + LlmRegistry.registerLlm("test-sim-.*", SIMULATION_MODELS::get); + } + + @Rule public final TestName testName = new TestName(); + + private String modelName; + + @Before + public void setUp() { + modelName = "test-sim-" + testName.getMethodName(); + } + + @Test + public void simulate_toolNotInTheConfiguration_letsTheToolRun() { + StubTool calendar = new StubTool("calendar"); + EnvironmentSimulationEngine engine = new EnvironmentSimulationEngine(mockingConfig("weather")); + + Map answer = + engine.simulate(calendar, ImmutableMap.of(), toolContextFor(calendar)).blockingGet(); + + assertThat(answer).isNull(); + } + + @Test + public void simulate_injectedError_answersWithTheError() { + StubTool weather = new StubTool("weather"); + EnvironmentSimulationEngine engine = + new EnvironmentSimulationEngine( + configOf( + ToolSimulationConfig.builder() + .toolName("weather") + .injectionConfigs( + ImmutableList.of( + InjectionConfig.builder() + .injectedError( + InjectedError.builder() + .injectedHttpErrorCode(503) + .errorMessage("upstream is down") + .build()) + .build())) + .build())); + + Map answer = + engine + .simulate(weather, ImmutableMap.of("city", "Zurich"), toolContextFor(weather)) + .blockingGet(); + + assertThat(answer).containsExactly("error_code", 503, "error_message", "upstream is down"); + } + + @Test + public void simulate_matchArgsAreASubsetOfTheCall_injects() { + StubTool weather = new StubTool("weather"); + EnvironmentSimulationEngine engine = + new EnvironmentSimulationEngine(matchingConfig(ImmutableMap.of("city", "Zurich"))); + + Map answer = + engine + .simulate( + weather, + ImmutableMap.of("city", "Zurich", "units", "celsius"), + toolContextFor(weather)) + .blockingGet(); + + assertThat(answer).containsExactly("temperature", 20); + } + + @Test + public void simulate_matchArgsDoNotMatchTheCall_letsTheToolRun() { + StubTool weather = new StubTool("weather"); + EnvironmentSimulationEngine engine = + new EnvironmentSimulationEngine(matchingConfig(ImmutableMap.of("city", "Zurich"))); + + Map answer = + engine + .simulate(weather, ImmutableMap.of("city", "Geneva"), toolContextFor(weather)) + .blockingGet(); + + assertThat(answer).isNull(); + } + + @Test + public void simulate_injectionProbabilityIsZero_letsTheToolRun() { + StubTool weather = new StubTool("weather"); + EnvironmentSimulationEngine engine = + new EnvironmentSimulationEngine( + configOf( + ToolSimulationConfig.builder() + .toolName("weather") + .injectionConfigs( + ImmutableList.of( + InjectionConfig.builder() + .injectionProbability(0.0) + .injectedResponse(ImmutableMap.of("temperature", 20)) + .build())) + .build())); + + Map answer = + engine.simulate(weather, ImmutableMap.of(), toolContextFor(weather)).blockingGet(); + + assertThat(answer).isNull(); + } + + @Test + public void simulate_seededInjection_drawsTheSameWayEveryRun() { + StubTool weather = new StubTool("weather"); + ToolContext toolContext = toolContextFor(weather); + EnvironmentSimulationConfig config = + configOf( + ToolSimulationConfig.builder() + .toolName("weather") + .injectionConfigs( + ImmutableList.of( + InjectionConfig.builder() + .randomSeed(42L) + .injectionProbability(0.5) + .injectedResponse(ImmutableMap.of("temperature", 20)) + .build())) + .build()); + + // Every call re-seeds before it draws, so all of them decide the same way. An engine that + // dropped the seed would draw around this threshold roughly half the time instead. + boolean expectedToInject = new Random(42L).nextDouble() < 0.5; + EnvironmentSimulationEngine engine = new EnvironmentSimulationEngine(config); + for (int call = 0; call < 10; call++) { + Map answer = + engine.simulate(weather, ImmutableMap.of(), toolContext).blockingGet(); + + assertThat(answer != null).isEqualTo(expectedToInject); + } + } + + @Test + public void simulate_mockStrategy_answersWithWhatTheModelWrote() { + answerWith(MOCK_RESPONSE_JSON); + StubTool weather = new StubTool("weather"); + EnvironmentSimulationEngine engine = new EnvironmentSimulationEngine(mockingConfig("weather")); + + Map answer = + engine.simulate(weather, ImmutableMap.of(), toolContextFor(weather)).blockingGet(); + + assertThat(answer).containsExactly("status", "ok"); + } + + @Test + public void simulate_repeatedCalls_analyzeToolConnectionsOnce() { + answerWith(MOCK_RESPONSE_JSON); + StubTool weather = new StubTool("weather"); + ToolContext toolContext = toolContextFor(weather); + EnvironmentSimulationEngine engine = new EnvironmentSimulationEngine(mockingConfig("weather")); + + for (int call = 0; call < 3; call++) { + Map unused = + engine.simulate(weather, ImmutableMap.of(), toolContext).blockingGet(); + } + + assertThat(promptsContaining(ANALYSIS_PROMPT)).isEqualTo(1); + assertThat(promptsContaining(MOCK_PROMPT)).isEqualTo(3); + } + + @Test + public void simulate_concurrentCalls_analyzeToolConnectionsOnce() throws Exception { + // The analysis is slow enough that every caller is inside simulate() while the first one is + // still waiting on the model. Reading a flag and setting it once the analysis returned would + // let each of them start an analysis of its own and overwrite the others' connection map. + answerAfter(MOCK_RESPONSE_JSON, Duration.ofMillis(300)); + StubTool weather = new StubTool("weather"); + ToolContext toolContext = toolContextFor(weather); + EnvironmentSimulationEngine engine = new EnvironmentSimulationEngine(mockingConfig("weather")); + + int callers = 8; + CyclicBarrier allReady = new CyclicBarrier(callers); + ExecutorService callerPool = Executors.newFixedThreadPool(callers); + List>> answers = new ArrayList<>(); + try { + for (int caller = 0; caller < callers; caller++) { + answers.add( + callerPool.submit( + () -> { + int unused = allReady.await(); + return engine.simulate(weather, ImmutableMap.of(), toolContext).blockingGet(); + })); + } + for (Future> answer : answers) { + assertThat(answer.get(60, SECONDS)).containsExactly("status", "ok"); + } + } finally { + callerPool.shutdownNow(); + } + + assertThat(promptsContaining(ANALYSIS_PROMPT)).isEqualTo(1); + assertThat(promptsContaining(MOCK_PROMPT)).isEqualTo(callers); + } + + @Test + public void simulate_toolConnectionAnalysisFails_stillAnswersWithAMock() { + // The first model call is the analysis and it fails; the second is the mock response. + registerSimulationModel( + createTestLlm( + Flowable.error(new IllegalStateException("model unavailable")), + Flowable.just(createTextLlmResponse(MOCK_RESPONSE_JSON)))); + StubTool weather = new StubTool("weather"); + EnvironmentSimulationEngine engine = new EnvironmentSimulationEngine(mockingConfig("weather")); + + Map answer = + engine.simulate(weather, ImmutableMap.of(), toolContextFor(weather)).blockingGet(); + + assertThat(answer).containsExactly("status", "ok"); + } + + private EnvironmentSimulationConfig configOf(ToolSimulationConfig... toolSimulationConfigs) { + return EnvironmentSimulationConfig.builder() + .simulationModel(modelName) + .toolSimulationConfigs(ImmutableList.copyOf(toolSimulationConfigs)) + .build(); + } + + private EnvironmentSimulationConfig mockingConfig(String toolName) { + return configOf( + ToolSimulationConfig.builder() + .toolName(toolName) + .mockStrategyType(MockStrategy.MOCK_STRATEGY_TOOL_SPEC) + .build()); + } + + private EnvironmentSimulationConfig matchingConfig(Map matchArgs) { + return configOf( + ToolSimulationConfig.builder() + .toolName("weather") + .injectionConfigs( + ImmutableList.of( + InjectionConfig.builder() + .matchArgs(matchArgs) + .injectedResponse(ImmutableMap.of("temperature", 20)) + .build())) + .build()); + } + + private ToolContext toolContextFor(BaseTool... tools) { + LlmAgent agent = + createTestAgentBuilder(new TestLlm(ImmutableList.of())) + .tools(ImmutableList.copyOf(tools)) + .build(); + return ToolContext.builder(createInvocationContext(agent)).build(); + } + + /** Points this test's simulation model at a fake that always answers with {@code json}. */ + private void answerWith(String json) { + registerSimulationModel(createTestLlm(() -> Flowable.just(createTextLlmResponse(json)))); + } + + /** The same, but the fake takes {@code latency} to answer. */ + private void answerAfter(String json, Duration latency) { + registerSimulationModel( + createTestLlm( + () -> + Flowable.just(createTextLlmResponse(json)) + .delay(latency.toMillis(), MILLISECONDS))); + } + + private void registerSimulationModel(TestLlm simulationModel) { + SIMULATION_MODELS.put(modelName, simulationModel); + } + + /** How many prompts the simulation sent to the model that contain {@code phrase}. */ + private long promptsContaining(String phrase) { + return SIMULATION_MODELS.get(modelName).getRequests().stream() + .map( + request -> + request.contents().stream() + .flatMap(content -> content.parts().orElse(ImmutableList.of()).stream()) + .map(part -> part.text().orElse("")) + .collect(joining())) + .filter(prompt -> prompt.contains(phrase)) + .count(); + } + + /** A tool that is never really called, and only has to have a name and a declaration. */ + private static final class StubTool extends BaseTool { + + StubTool(String name) { + super(name, "the " + name + " tool"); + } + + @Override + public Optional declaration() { + return Optional.of( + FunctionDeclaration.builder().name(name()).description(description()).build()); + } + } +}