diff --git a/core/src/main/java/com/google/adk/agents/ContextCacheConfig.java b/core/src/main/java/com/google/adk/agents/ContextCacheConfig.java index 084700d54..571fc15e3 100644 --- a/core/src/main/java/com/google/adk/agents/ContextCacheConfig.java +++ b/core/src/main/java/com/google/adk/agents/ContextCacheConfig.java @@ -15,45 +15,90 @@ */ package com.google.adk.agents; +import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.base.Preconditions.checkNotNull; + +import com.google.errorprone.annotations.InlineMe; +import com.google.genai.types.HttpOptions; import java.time.Duration; +import org.jspecify.annotations.Nullable; /** - * Configuration for context caching across all agents in an app. - * - *

This configuration enables and controls context caching behavior for all LLM agents in an app. - * When this config is present on an app, context caching is enabled for all agents. When absent - * (null), context caching is disabled. - * - *

Context caching can significantly reduce costs and improve response times by reusing - * previously processed context across multiple requests. + * Configuration for context caching across all agents in an app; without it, nothing is cached. + * Gemini models cache the stable prefix of an agent's requests from its second request on, once + * that prefix reaches the model's minimum: 2048 tokens for Gemini 2.5, 4096 for Gemini 3. Other + * models ignore this config, and a cache carries over to the next invocation only with a session + * service that stores whole events, such as {@code InMemorySessionService}. * - * @param maxInvocations Maximum number of invocations to reuse the same cache before refreshing it. - * Defaults to 10. - * @param ttl Time-to-live for cache. Defaults to 1800 seconds (30 minutes). - * @param minTokens Minimum estimated request tokens required to enable caching. This compares - * against the estimated total tokens of the request (system instruction + tools + contents). - * Context cache storage may have cost. Set higher to avoid caching small requests where - * overhead may exceed benefits. Defaults to 0. + * @param cacheIntervals Maximum number of invocations to reuse the same cache before refreshing it, + * from 1 to 100. Defaults to 10. + * @param ttl Time-to-live for cache; a whole number of seconds, at least one. Defaults to 1800 + * seconds (30 minutes). + * @param minTokens Minimum prompt token count of the agent's previous request needed to create a + * cache; raise it to skip small requests, where cache storage can cost more than it saves. Must + * not be negative. Defaults to 0. + * @param createHttpOptions HTTP options, such as a timeout, for the call that creates a cache; null + * uses the client's defaults. If that call fails, the request is sent without a cache. Defaults + * to null. */ -public record ContextCacheConfig(int maxInvocations, Duration ttl, int minTokens) { +public record ContextCacheConfig( + int cacheIntervals, Duration ttl, int minTokens, @Nullable HttpOptions createHttpOptions) { + + /** + * Validates the config as ADK Python does. + * + * @throws IllegalArgumentException if a value is out of range + * @throws NullPointerException if {@code ttl} is null + */ + public ContextCacheConfig { + checkArgument( + cacheIntervals >= 1 && cacheIntervals <= 100, + "cacheIntervals must be between 1 and 100, but was %s.", + cacheIntervals); + checkNotNull(ttl, "ttl must not be null."); + checkArgument( + ttl.compareTo(Duration.ofSeconds(1)) >= 0 && ttl.toNanosPart() == 0, + "ttl must be a whole number of seconds, at least one, but was %s.", + ttl); + checkArgument(minTokens >= 0, "minTokens must not be negative, but was %s.", minTokens); + } public ContextCacheConfig() { - this(10, Duration.ofSeconds(1800), 0); + this(10, Duration.ofMinutes(30), 0); + } + + /** Creates a config that creates caches with the client's default HTTP options. */ + public ContextCacheConfig(int cacheIntervals, Duration ttl, int minTokens) { + this(cacheIntervals, ttl, minTokens, /* createHttpOptions= */ null); + } + + /** + * Returns {@link #cacheIntervals()}. + * + * @deprecated Use {@link #cacheIntervals()}, the name ADK Python and ADK Kotlin use. + */ + @Deprecated + @InlineMe(replacement = "this.cacheIntervals()") + public int maxInvocations() { + return cacheIntervals(); } - /** Returns TTL as string format for cache creation. */ + /** Returns the TTL in the {@code "s"} form that the Gemini API uses. */ public String getTtlString() { - return ttl.getSeconds() + "s"; + return ttl.toSeconds() + "s"; } @Override public String toString() { - return "ContextCacheConfig(maxInvocations=" - + maxInvocations + // Says only whether HTTP options are set, since their headers can carry credentials. + return "ContextCacheConfig(cacheIntervals=" + + cacheIntervals + ", ttl=" - + ttl.getSeconds() + + ttl.toSeconds() + "s, minTokens=" + minTokens + + ", createHttpOptions=" + + (createHttpOptions == null ? "null" : "set") + ")"; } } diff --git a/core/src/main/java/com/google/adk/events/Event.java b/core/src/main/java/com/google/adk/events/Event.java index 5b63be76f..eb266f459 100644 --- a/core/src/main/java/com/google/adk/events/Event.java +++ b/core/src/main/java/com/google/adk/events/Event.java @@ -24,6 +24,7 @@ import com.fasterxml.jackson.databind.annotation.JsonDeserialize; import com.google.adk.JsonBaseModel; import com.google.adk.annotations.Experimental; +import com.google.adk.models.CacheMetadata; import com.google.adk.platform.UuidProvider; import com.google.adk.workflow.NodeInfo; import com.google.common.collect.ImmutableList; @@ -70,6 +71,7 @@ public class Event extends JsonBaseModel { private @Nullable Transcription outputTranscription; private @Nullable Object output; private @Nullable NodeInfo nodeInfo; + private @Nullable CacheMetadata cacheMetadata; private long timestamp; @@ -342,6 +344,19 @@ public void setNodeInfo(@Nullable NodeInfo nodeInfo) { this.nodeInfo = nodeInfo; } + /** + * Context cache state of the LLM response this event carries. The next request of the same agent + * reads it to reuse or refresh the cache. + */ + @JsonProperty("cacheMetadata") + public Optional cacheMetadata() { + return Optional.ofNullable(cacheMetadata); + } + + public void setCacheMetadata(@Nullable CacheMetadata cacheMetadata) { + this.cacheMetadata = cacheMetadata; + } + /** The timestamp of the event. */ @JsonProperty("timestamp") public long timestamp() { @@ -453,6 +468,7 @@ public static class Builder { private @Nullable Transcription outputTranscription; private @Nullable Object output; private @Nullable NodeInfo nodeInfo; + private @Nullable CacheMetadata cacheMetadata; private @Nullable Long timestamp; @JsonCreator @@ -646,6 +662,13 @@ public Builder nodeInfo(@Nullable NodeInfo value) { return this; } + @CanIgnoreReturnValue + @JsonProperty("cacheMetadata") + public Builder cacheMetadata(@Nullable CacheMetadata value) { + this.cacheMetadata = value; + return this; + } + public Event build() { Event event = new Event(); event.setId(id); @@ -672,6 +695,7 @@ public Event build() { event.setOutputTranscription(outputTranscription); event.setOutput(output); event.setNodeInfo(nodeInfo); + event.setCacheMetadata(cacheMetadata); return event; } } @@ -710,6 +734,7 @@ public Builder toBuilder() { .outputTranscription(this.outputTranscription) .output(this.output) .nodeInfo(this.nodeInfo) + .cacheMetadata(this.cacheMetadata) .timestamp(this.timestamp); } @@ -743,7 +768,8 @@ public boolean equals(Object obj) { && Objects.equals(inputTranscription, other.inputTranscription) && Objects.equals(outputTranscription, other.outputTranscription) && Objects.equals(output, other.output) - && Objects.equals(nodeInfo, other.nodeInfo); + && Objects.equals(nodeInfo, other.nodeInfo) + && Objects.equals(cacheMetadata, other.cacheMetadata); } @Override @@ -776,6 +802,7 @@ public int hashCode() { outputTranscription, output, nodeInfo, + cacheMetadata, timestamp); } } 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 b576a299e..2b921ab22 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 @@ -885,7 +885,8 @@ private Event buildModelResponseEvent( .usageMetadata(llmResponse.usageMetadata().orElse(null)) .modelVersion(llmResponse.modelVersion().orElse(null)) .inputTranscription(llmResponse.inputTranscription().orElse(null)) - .outputTranscription(llmResponse.outputTranscription().orElse(null)); + .outputTranscription(llmResponse.outputTranscription().orElse(null)) + .cacheMetadata(llmResponse.cacheMetadata().orElse(null)); Event event = eventBuilder.build(); diff --git a/core/src/main/java/com/google/adk/flows/llmflows/ContextCacheRequestProcessor.java b/core/src/main/java/com/google/adk/flows/llmflows/ContextCacheRequestProcessor.java new file mode 100644 index 000000000..bfef398f1 --- /dev/null +++ b/core/src/main/java/com/google/adk/flows/llmflows/ContextCacheRequestProcessor.java @@ -0,0 +1,101 @@ +/* + * 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.flows.llmflows; + +import com.google.adk.agents.ContextCacheConfig; +import com.google.adk.agents.InvocationContext; +import com.google.adk.events.Event; +import com.google.adk.models.CacheMetadata; +import com.google.adk.models.LlmRequest; +import com.google.common.base.Strings; +import com.google.common.collect.ImmutableList; +import com.google.genai.types.GenerateContentResponseUsageMetadata; +import io.reactivex.rxjava3.core.Single; +import java.util.Optional; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * {@link RequestProcessor} that enables context caching when the app configures it. It puts the + * config, the agent's latest cache metadata and its previous prompt token count on the request; the + * model creates, reuses and deletes the caches. + */ +final class ContextCacheRequestProcessor implements RequestProcessor { + + private static final Logger logger = LoggerFactory.getLogger(ContextCacheRequestProcessor.class); + + @Override + public Single processRequest( + InvocationContext context, LlmRequest request) { + Optional cacheConfig = context.contextCacheConfig(); + if (cacheConfig.isEmpty()) { + return Single.just(RequestProcessingResult.create(request, ImmutableList.of())); + } + + String agentName = context.agent().name(); + CacheMetadata cacheMetadata = null; + Integer previousTokenCount = null; + for (Event event : context.session().immutableEvents().reverse()) { + if (!agentName.equals(event.author())) { + continue; + } + if (cacheMetadata == null && event.cacheMetadata().isPresent()) { + cacheMetadata = countInvocation(event, context.invocationId()); + } + if (previousTokenCount == null) { + previousTokenCount = + event + .usageMetadata() + .flatMap(GenerateContentResponseUsageMetadata::promptTokenCount) + .orElse(null); + } + if (cacheMetadata != null && previousTokenCount != null) { + break; + } + } + if (cacheMetadata != null) { + logger.debug("Found cache metadata for agent {}: {}", agentName, cacheMetadata); + } + if (previousTokenCount != null) { + logger.debug( + "Found previous prompt token count for agent {}: {}", agentName, previousTokenCount); + } + logger.debug("Context caching enabled for agent {}", agentName); + + LlmRequest updatedRequest = + request.toBuilder() + .cacheConfig(cacheConfig.get()) + .cacheMetadata(cacheMetadata) + .cacheableContentsTokenCount(previousTokenCount) + .build(); + return Single.just(RequestProcessingResult.create(updatedRequest, ImmutableList.of())); + } + + /** + * Returns the event's cache metadata, counting one more use when it names an active cache from an + * earlier invocation. + */ + private static CacheMetadata countInvocation(Event event, String invocationId) { + CacheMetadata metadata = event.cacheMetadata().get(); + if (Strings.isNullOrEmpty(event.invocationId()) + || event.invocationId().equals(invocationId) + || metadata.cacheName().isEmpty()) { + return metadata; + } + return metadata.toBuilder().invocationsUsed(metadata.invocationsUsed().get() + 1).build(); + } +} diff --git a/core/src/main/java/com/google/adk/flows/llmflows/SingleFlow.java b/core/src/main/java/com/google/adk/flows/llmflows/SingleFlow.java index 41dff3b96..cf4382e28 100644 --- a/core/src/main/java/com/google/adk/flows/llmflows/SingleFlow.java +++ b/core/src/main/java/com/google/adk/flows/llmflows/SingleFlow.java @@ -33,6 +33,7 @@ public class SingleFlow extends BaseLlmFlow { new Identity(), new Compaction(), new Contents(), + new ContextCacheRequestProcessor(), CodeExecution.requestProcessor); protected static final ImmutableList RESPONSE_PROCESSORS = diff --git a/core/src/main/java/com/google/adk/models/CacheMetadata.java b/core/src/main/java/com/google/adk/models/CacheMetadata.java new file mode 100644 index 000000000..576714a4b --- /dev/null +++ b/core/src/main/java/com/google/adk/models/CacheMetadata.java @@ -0,0 +1,202 @@ +/* + * 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.models; + +import static com.google.common.base.Preconditions.checkState; + +import com.fasterxml.jackson.annotation.JsonAlias; +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.fasterxml.jackson.core.JsonParser; +import com.fasterxml.jackson.core.JsonToken; +import com.fasterxml.jackson.databind.DeserializationContext; +import com.fasterxml.jackson.databind.JsonDeserializer; +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 java.io.IOException; +import java.math.BigDecimal; +import java.time.Duration; +import java.time.Instant; +import java.time.format.DateTimeParseException; +import java.util.Locale; +import java.util.Optional; +import org.jspecify.annotations.Nullable; + +/** + * Context cache state carried on an LLM response, and on its event, from one request to the next. + * + *

Metadata for an active cache has a {@link #cacheName()}, {@link #expireTime()} and {@link + * #invocationsUsed()}. Fingerprint-only metadata has none of them: it records the cacheable prefix + * so that a later request can tell whether the prefix is still the same. + */ +@AutoValue +@JsonDeserialize(builder = CacheMetadata.Builder.class) +public abstract class CacheMetadata extends JsonBaseModel { + + CacheMetadata() {} + + /** Hash of the cacheable state: system instruction, tools and the cached leading contents. */ + @JsonProperty("fingerprint") + public abstract String fingerprint(); + + /** Number of leading contents in the cache, or in the fingerprint when no cache is active. */ + @JsonProperty("contentsCount") + public abstract int contentsCount(); + + /** Full resource name of the cached content; empty when no cache is active. */ + @JsonProperty("cacheName") + public abstract Optional cacheName(); + + /** When the cache expires; empty when no cache is active. */ + @JsonProperty("expireTime") + public abstract Optional expireTime(); + + /** Number of invocations that have used the cache; empty when no cache is active. */ + @JsonProperty("invocationsUsed") + public abstract Optional invocationsUsed(); + + /** When the cache was created; empty when no cache is active. */ + @JsonProperty("createdAt") + public abstract Optional createdAt(); + + public abstract Builder toBuilder(); + + public static Builder builder() { + return new AutoValue_CacheMetadata.Builder(); + } + + /** Returns a short description for logs, in the same shape as ADK Python's. */ + @Override + public final String toString() { + if (cacheName().isEmpty()) { + String shortFingerprint = fingerprint().substring(0, Math.min(8, fingerprint().length())); + return "Fingerprint-only: " + + contentsCount() + + " contents, fingerprint=" + + shortFingerprint + + "..."; + } + String cacheId = cacheName().get().substring(cacheName().get().lastIndexOf('/') + 1); + double minutesToExpiry = + Duration.between(Instant.now(), expireTime().get()).toMillis() / 60_000.0; + return String.format( + Locale.ROOT, + "Cache %s: used %d invocations, cached %d contents, expires in %.1fmin", + cacheId, + invocationsUsed().get(), + contentsCount(), + minutesToExpiry); + } + + /** + * Builder for {@link CacheMetadata}. The snake_case aliases read metadata that ADK Python wrote. + */ + @AutoValue.Builder + @JsonPOJOBuilder(buildMethodName = "build", withPrefix = "") + public abstract static class Builder { + + @JsonCreator + static Builder jacksonBuilder() { + return CacheMetadata.builder(); + } + + @JsonProperty("fingerprint") + public abstract Builder fingerprint(String fingerprint); + + @JsonProperty("contentsCount") + @JsonAlias("contents_count") + public abstract Builder contentsCount(int contentsCount); + + @JsonProperty("cacheName") + @JsonAlias("cache_name") + public abstract Builder cacheName(@Nullable String cacheName); + + @JsonProperty("expireTime") + @JsonAlias("expire_time") + @JsonDeserialize(using = LenientEpochDeserializer.class) + public abstract Builder expireTime(@Nullable Instant expireTime); + + @JsonProperty("invocationsUsed") + @JsonAlias("invocations_used") + public abstract Builder invocationsUsed(@Nullable Integer invocationsUsed); + + @JsonProperty("createdAt") + @JsonAlias("created_at") + @JsonDeserialize(using = LenientEpochDeserializer.class) + public abstract Builder createdAt(@Nullable Instant createdAt); + + abstract CacheMetadata autoBuild(); + + /** + * Builds the metadata. + * + * @throws IllegalStateException if a count is negative, or if only some of {@code cacheName}, + * {@code expireTime} and {@code invocationsUsed} are set + */ + public CacheMetadata build() { + CacheMetadata metadata = autoBuild(); + checkState(metadata.contentsCount() >= 0, "contentsCount must not be negative."); + checkState( + metadata.invocationsUsed().orElse(0) >= 0, "invocationsUsed must not be negative."); + boolean active = metadata.cacheName().isPresent(); + checkState( + metadata.expireTime().isPresent() == active + && metadata.invocationsUsed().isPresent() == active, + "cacheName, expireTime and invocationsUsed must be all set or all unset."); + return metadata; + } + } + + /** + * Reads a numeric timestamp as epoch seconds when its magnitude is below 1e11 and as epoch + * milliseconds otherwise, since ADK Python writes seconds and ADK Kotlin writes milliseconds. + * Also reads ISO-8601 text. + */ + static final class LenientEpochDeserializer extends JsonDeserializer { + private static final BigDecimal SECONDS_MILLIS_BOUNDARY = new BigDecimal("1e11"); + // Years 1 to 9999, as ADK Kotlin accepts; checked before any arithmetic that could be costly. + private static final BigDecimal MIN_SECONDS = BigDecimal.valueOf(-62_135_596_800L); + private static final BigDecimal MAX_SECONDS = BigDecimal.valueOf(253_402_300_799L); + + @Override + public Instant deserialize(JsonParser parser, DeserializationContext context) + throws IOException { + if (parser.currentToken() == JsonToken.VALUE_STRING) { + try { + return Instant.parse(parser.getText()); + } catch (DateTimeParseException e) { + return (Instant) + context.handleWeirdStringValue( + Instant.class, parser.getText(), "not an ISO-8601 time"); + } + } + if (!parser.currentToken().isNumeric()) { + return (Instant) context.handleUnexpectedToken(Instant.class, parser); + } + BigDecimal value = parser.getDecimalValue(); + BigDecimal seconds = + value.abs().compareTo(SECONDS_MILLIS_BOUNDARY) < 0 ? value : value.movePointLeft(3); + if (seconds.compareTo(MIN_SECONDS) < 0 || seconds.compareTo(MAX_SECONDS) > 0) { + return (Instant) context.handleWeirdNumberValue(Instant.class, value, "out of range"); + } + return Instant.ofEpochSecond( + seconds.longValue(), seconds.remainder(BigDecimal.ONE).movePointRight(9).longValue()); + } + } +} diff --git a/core/src/main/java/com/google/adk/models/Gemini.java b/core/src/main/java/com/google/adk/models/Gemini.java index ca6c89550..e4c3c0c25 100644 --- a/core/src/main/java/com/google/adk/models/Gemini.java +++ b/core/src/main/java/com/google/adk/models/Gemini.java @@ -273,11 +273,26 @@ public Gemini build() { @Override public Flowable generateContent(LlmRequest llmRequest, boolean stream) { - llmRequest = + LlmRequest preparedRequest = GeminiUtil.prepareGenenerateContentRequest( llmRequest, !apiClient.vertexAI(), /* stripThoughts= */ false); + String effectiveModelName = preparedRequest.model().orElse(model()); + if (preparedRequest.cacheConfig().isEmpty()) { + return sendRequest(effectiveModelName, preparedRequest, stream, /* cacheMetadata= */ null); + } + return new GeminiContextCacheManager(apiClient, effectiveModelName) + .handleContextCaching(preparedRequest) + .flatMapPublisher( + cached -> sendRequest(effectiveModelName, cached.request(), stream, cached.metadata())); + } + + /** Sends the request, and puts {@code cacheMetadata}, if any, on the final response. */ + private Flowable sendRequest( + String effectiveModelName, + LlmRequest llmRequest, + boolean stream, + @Nullable CacheMetadata cacheMetadata) { GenerateContentConfig config = llmRequest.config().orElse(null); - String effectiveModelName = llmRequest.model().orElse(model()); logger.trace("Request Contents: {}", llmRequest.contents()); logger.trace("Request Config: {}", config); @@ -290,18 +305,29 @@ public Flowable generateContent(LlmRequest llmRequest, boolean stre return Flowable.defer( () -> streamResumingPauses( - effectiveModelName, new GeminiContinuation(contents, config), streamFuture)); + effectiveModelName, + new GeminiContinuation(contents, config), + streamFuture, + cacheMetadata)); } else { logger.debug("Sending generateContent request to model {}", effectiveModelName); CompletableFuture responseFuture = apiClient.async.models.generateContent(effectiveModelName, contents, config); return Flowable.defer( - () -> - generateResumingPauses( - effectiveModelName, new GeminiContinuation(contents, config), responseFuture)); + () -> + generateResumingPauses( + effectiveModelName, new GeminiContinuation(contents, config), responseFuture)) + .map(response -> withCacheMetadata(response, cacheMetadata)); } } + private static LlmResponse withCacheMetadata( + LlmResponse response, @Nullable CacheMetadata cacheMetadata) { + return cacheMetadata == null + ? response + : response.toBuilder().cacheMetadata(cacheMetadata).build(); + } + @VisibleForTesting static Flowable processRawResponses(Flowable rawResponses) { return Flowable.defer(() -> new StreamingResponseAggregator().process(rawResponses)); @@ -353,7 +379,8 @@ private Optional completeOrResume( private Flowable streamResumingPauses( String modelName, GeminiContinuation continuation, - CompletableFuture> firstStream) { + CompletableFuture> firstStream, + @Nullable CacheMetadata cacheMetadata) { Deque>> streams = new ArrayDeque<>(); streams.add(firstStream); StreamingResponseAggregator aggregator = new StreamingResponseAggregator(); @@ -362,7 +389,11 @@ private Flowable streamResumingPauses( .concatMap(aggregator::processRawResponse) .concatWith( Flowable.defer( - () -> aggregator.processFinalResponse().map(continuation::withSummedUsage))); + () -> + aggregator + .processFinalResponse() + .map(continuation::withSummedUsage) + .map(response -> withCacheMetadata(response, cacheMetadata)))); } /** diff --git a/core/src/main/java/com/google/adk/models/GeminiContextCacheManager.java b/core/src/main/java/com/google/adk/models/GeminiContextCacheManager.java new file mode 100644 index 000000000..5cf8924fd --- /dev/null +++ b/core/src/main/java/com/google/adk/models/GeminiContextCacheManager.java @@ -0,0 +1,485 @@ +/* + * 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.models; + +import static java.nio.charset.StandardCharsets.UTF_8; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.node.ArrayNode; +import com.fasterxml.jackson.databind.node.JsonNodeFactory; +import com.fasterxml.jackson.databind.node.ObjectNode; +import com.google.adk.agents.ContextCacheConfig; +import com.google.adk.agents.Role; +import com.google.common.annotations.VisibleForTesting; +import com.google.common.base.Ascii; +import com.google.common.collect.ImmutableList; +import com.google.genai.Client; +import com.google.genai.JsonSerializable; +import com.google.genai.errors.ApiException; +import com.google.genai.types.CachedContent; +import com.google.genai.types.Content; +import com.google.genai.types.CreateCachedContentConfig; +import com.google.genai.types.GenerateContentConfig; +import com.google.genai.types.HttpOptions; +import com.google.genai.types.Tool; +import io.reactivex.rxjava3.core.Completable; +import io.reactivex.rxjava3.core.Maybe; +import io.reactivex.rxjava3.core.Single; +import java.security.MessageDigest; +import java.security.NoSuchAlgorithmException; +import java.time.Instant; +import java.time.InstantSource; +import java.util.ArrayList; +import java.util.Comparator; +import java.util.HexFormat; +import java.util.List; +import java.util.Optional; +import java.util.OptionalInt; +import java.util.concurrent.ExecutionException; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Creates, reuses and deletes the Gemini context cache of a request that carries a {@link + * ContextCacheConfig}. A cache covers the system instruction, tools and leading contents, and a + * fingerprint of them decides whether a later request can reuse it. + */ +final class GeminiContextCacheManager { + + private static final Logger logger = LoggerFactory.getLogger(GeminiContextCacheManager.class); + + // Documented explicit-cache minimums; for tuned-model and endpoint IDs the server decides. + private static final int GEMINI_2_5_MIN_CACHE_TOKENS = 2048; + private static final int GEMINI_3_MIN_CACHE_TOKENS = 4096; + + private static final int FINGERPRINT_LENGTH = 16; + private static final int CHARS_PER_TOKEN = 4; + private static final JsonNodeFactory NODES = JsonNodeFactory.instance; + + private final Client apiClient; + private final String model; + private final InstantSource clock; + + /** Creates a manager for caches of {@code model}, the model the requests are sent to. */ + GeminiContextCacheManager(Client apiClient, String model) { + this(apiClient, model, InstantSource.system()); + } + + /** Creates a manager that reads the current time from {@code clock}, for tests. */ + @VisibleForTesting + GeminiContextCacheManager(Client apiClient, String model, InstantSource clock) { + this.apiClient = apiClient; + this.model = model; + this.clock = clock; + } + + /** + * The request to send, which uses a cache when one applies, and the metadata for its response. + */ + record CacheResult(LlmRequest request, CacheMetadata metadata) {} + + /** + * Reuses the request's cache while it is valid, and otherwise deletes it and creates a new one + * once the request's cacheable prefix is known to be stable. The returned single fails with + * {@link IllegalArgumentException} if the request has no cache config. + */ + Single handleContextCaching(LlmRequest request) { + return Single.defer( + () -> { + ContextCacheConfig config = + request + .cacheConfig() + .orElseThrow( + () -> + new IllegalArgumentException( + "Context caching requires a cache configuration.")); + Optional previous = request.cacheMetadata(); + if (previous.isEmpty()) { + // A cache is never created without an earlier fingerprint to compare with. + return Single.just(new CacheResult(request, fingerprintOnly(request))); + } + CacheMetadata metadata = previous.get(); + if (isCacheValid(request, metadata, config)) { + logger.debug("Reusing cache {}.", metadata.cacheName().get()); + return Single.just( + new CacheResult( + applyCache(request, metadata.cacheName().get(), metadata.contentsCount()), + metadata)); + } + return metadata + .cacheName() + .map(this::deleteCache) + .orElseGet(Completable::complete) + .andThen(Single.defer(() -> refreshCache(request, metadata, config))); + }); + } + + /** + * Creates a cache when the previously fingerprinted prefix is unchanged, growing it to the + * current cacheable prefix; otherwise starts a new fingerprint chain. + */ + private Single refreshCache( + LlmRequest request, CacheMetadata previous, ContextCacheConfig config) { + if (!fingerprint(request, previous.contentsCount()).equals(previous.fingerprint())) { + logger.debug("Cacheable prefix changed; returning fingerprint-only metadata."); + return Single.just(new CacheResult(request, fingerprintOnly(request))); + } + int contentsCount = + Math.max(previous.contentsCount(), countContentsToCache(request.contents())); + // Keeping the grown prefix when creation fails lets the next request try again. + CacheMetadata grownPrefix = fingerprintOnly(request, contentsCount); + return createCache(request, grownPrefix, config) + .map( + created -> + new CacheResult( + applyCache(request, created.cacheName().get(), contentsCount), created)) + .defaultIfEmpty(new CacheResult(request, grownPrefix)); + } + + private CacheMetadata fingerprintOnly(LlmRequest request) { + return fingerprintOnly(request, countContentsToCache(request.contents())); + } + + private CacheMetadata fingerprintOnly(LlmRequest request, int contentsCount) { + return CacheMetadata.builder() + .fingerprint(fingerprint(request, contentsCount)) + .contentsCount(contentsCount) + .build(); + } + + /** Returns the number of contents before the last run of user contents. */ + private static int countContentsToCache(List contents) { + int lastUserRunStart = contents.size(); + for (int i = contents.size() - 1; i >= 0; i--) { + if (!Ascii.equalsIgnoreCase(contents.get(i).role().orElse(""), Role.USER)) { + break; + } + lastUserRunStart = i; + } + return lastUserRunStart; + } + + private boolean isCacheValid( + LlmRequest request, CacheMetadata metadata, ContextCacheConfig config) { + if (metadata.cacheName().isEmpty()) { + return false; + } + String cacheName = metadata.cacheName().get(); + if (!clock.instant().isBefore(metadata.expireTime().get())) { + logger.info("Cache {} expired.", cacheName); + return false; + } + if (metadata.invocationsUsed().get() > config.cacheIntervals()) { + logger.info( + "Cache {} was used by {} invocations, more than {}.", + cacheName, + metadata.invocationsUsed().get(), + config.cacheIntervals()); + return false; + } + if (!fingerprint(request, metadata.contentsCount()).equals(metadata.fingerprint())) { + logger.debug("Cache {} no longer matches the request.", cacheName); + return false; + } + return true; + } + + /** + * Returns a hash of everything a cache of the first {@code contentsCount} contents would hold, + * plus the model and backend that own the cache. Keys are sorted and tools are put in a fixed + * order, so equivalent requests always hash the same. + */ + private String fingerprint(LlmRequest request, int contentsCount) { + ObjectNode data = NODES.objectNode(); + data.put("model", model); + data.set("cache_scope", cacheScope()); + request + .config() + .ifPresent( + config -> { + config + .systemInstruction() + .ifPresent(instruction -> data.set("system_instruction", tree(instruction))); + config + .tools() + .filter(tools -> !tools.isEmpty()) + .ifPresent(tools -> data.set("tools", canonicalTools(tools))); + config + .toolConfig() + .ifPresent(toolConfig -> data.set("tool_config", tree(toolConfig))); + }); + List cachedContents = leading(request.contents(), contentsCount); + if (!cachedContents.isEmpty()) { + ArrayNode contents = data.putArray("cached_contents"); + cachedContents.forEach(content -> contents.add(tree(content))); + } + return sha256Hex(data.toString()).substring(0, FINGERPRINT_LENGTH); + } + + /** The backend namespace that owns the cache, so a cache is never reused across backends. */ + private ObjectNode cacheScope() { + ObjectNode scope = NODES.objectNode(); + scope.put("backend", apiClient.vertexAI() ? "vertex" : "gemini"); + if (apiClient.vertexAI()) { + scope.put("project", apiClient.project()); + scope.put("location", apiClient.location()); + } + return scope; + } + + private static ArrayNode canonicalTools(List tools) { + List canonical = new ArrayList<>(); + for (Tool tool : tools) { + ObjectNode node = (ObjectNode) tree(tool); + JsonNode declarations = node.get("functionDeclarations"); + if (declarations != null && declarations.isArray()) { + List sorted = new ArrayList<>(); + declarations.forEach(sorted::add); + sorted.sort(Comparator.comparing(declaration -> declaration.path("name").asText(""))); + node.set("functionDeclarations", NODES.arrayNode().addAll(sorted)); + } + canonical.add(node); + } + canonical.sort(Comparator.comparing(JsonNode::toString)); + return NODES.arrayNode().addAll(canonical); + } + + /** Returns the wire JSON of a GenAI type, with object keys sorted at every level. */ + private static JsonNode tree(Object genaiObject) { + return sortedKeys(JsonSerializable.toJsonNode(genaiObject)); + } + + private static JsonNode sortedKeys(JsonNode node) { + if (node.isObject()) { + ObjectNode sorted = NODES.objectNode(); + ImmutableList.sortedCopyOf(ImmutableList.copyOf(node.fieldNames())) + .forEach(name -> sorted.set(name, sortedKeys(node.get(name)))); + return sorted; + } + if (node.isArray()) { + ArrayNode sorted = NODES.arrayNode(); + node.forEach(element -> sorted.add(sortedKeys(element))); + return sorted; + } + return node; + } + + private static String sha256Hex(String text) { + try { + return HexFormat.of() + .formatHex(MessageDigest.getInstance("SHA-256").digest(text.getBytes(UTF_8))); + } catch (NoSuchAlgorithmException impossible) { + throw new AssertionError(impossible); + } + } + + /** + * Creates a cache of the system instruction, tools and the contents {@code prefix} counts, or + * completes empty when the request is too small to cache or creation fails. + */ + private Maybe createCache( + LlmRequest request, CacheMetadata prefix, ContextCacheConfig config) { + int contentsCount = prefix.contentsCount(); + Optional previousTokenCount = request.cacheableContentsTokenCount(); + if (previousTokenCount.isEmpty()) { + logger.info("No previous token count available; not creating a cache."); + return Maybe.empty(); + } + if (previousTokenCount.get() < config.minTokens()) { + logger.info( + "Previous request too small for caching ({} < {} tokens).", + previousTokenCount.get(), + config.minTokens()); + return Maybe.empty(); + } + OptionalInt minimumTokens = minimumCacheTokens(model); + if (minimumTokens.isPresent()) { + int prefixTokens = + estimateCacheablePrefixTokens(request, contentsCount, previousTokenCount.get()); + if (prefixTokens < minimumTokens.getAsInt()) { + logger.info( + "Cacheable prefix below the model's minimum cache size ({} < {} tokens).", + prefixTokens, + minimumTokens.getAsInt()); + return Maybe.empty(); + } + } + + CreateCachedContentConfig.Builder cacheConfig = + CreateCachedContentConfig.builder() + .ttl(config.ttl()) + .displayName( + "adk-cache-" + clock.instant().getEpochSecond() + "-" + contentsCount + "contents"); + HttpOptions createHttpOptions = config.createHttpOptions(); + if (createHttpOptions != null) { + cacheConfig.httpOptions(createHttpOptions); + } + List cachedContents = leading(request.contents(), contentsCount); + if (!cachedContents.isEmpty()) { + cacheConfig.contents(cachedContents); + } + request + .config() + .ifPresent( + generateConfig -> { + generateConfig.systemInstruction().ifPresent(cacheConfig::systemInstruction); + generateConfig + .tools() + .filter(tools -> !tools.isEmpty()) + .ifPresent(cacheConfig::tools); + generateConfig.toolConfig().ifPresent(cacheConfig::toolConfig); + }); + + // Blocking, as Gemini does on generateContent, keeps the model call on the caller's thread. + return Maybe.defer( + () -> Maybe.fromFuture(apiClient.async.caches.create(model, cacheConfig.build()))) + .map(cachedContent -> toMetadata(cachedContent, prefix, config)) + .doOnSuccess(metadata -> logger.info("Created cache {}.", metadata.cacheName().get())) + .onErrorResumeNext( + error -> { + if (error instanceof InterruptedException) { + Thread.currentThread().interrupt(); + return Maybe.error(error); + } + logger.warn("Failed to create cache: {}", describe(error)); + return Maybe.empty(); + }); + } + + private CacheMetadata toMetadata( + CachedContent cachedContent, CacheMetadata prefix, ContextCacheConfig config) { + Instant createdAt = clock.instant(); + String cacheName = + cachedContent + .name() + .filter(name -> !name.isEmpty()) + .orElseThrow( + () -> new IllegalStateException("The cache service returned no cache name.")); + return prefix.toBuilder() + .cacheName(cacheName) + .expireTime(cachedContent.expireTime().orElse(createdAt.plus(config.ttl()))) + .invocationsUsed(1) + .createdAt(createdAt) + .build(); + } + + private static OptionalInt minimumCacheTokens(String model) { + String name = model.substring(model.lastIndexOf('/') + 1); + if (name.startsWith("gemini-2.5-")) { + return OptionalInt.of(GEMINI_2_5_MIN_CACHE_TOKENS); + } + if (name.startsWith("gemini-3")) { + return OptionalInt.of(GEMINI_3_MIN_CACHE_TOKENS); + } + return OptionalInt.empty(); + } + + /** + * Scales the previous prompt's token count by the share of characters the cached prefix has in + * the request, since only the whole prompt has an exact count. + */ + private static int estimateCacheablePrefixTokens( + LlmRequest request, int contentsCount, int fullTokens) { + long fullEstimate = estimateTokens(request, request.contents().size()); + if (fullEstimate <= 0) { + // Nothing to estimate from, such as binary-only parts: trust the exact count. + return fullTokens; + } + long prefixEstimate = estimateTokens(request, contentsCount); + double ratio = Math.min(1.0, (double) prefixEstimate / fullEstimate); + return (int) (fullTokens * ratio); + } + + /** Roughly estimates the tokens of the system instruction, tools and leading contents. */ + private static long estimateTokens(LlmRequest request, int contentsCount) { + long chars = 0; + Optional config = request.config(); + if (config.isPresent()) { + chars += + config.get().systemInstruction().map(GeminiContextCacheManager::textLength).orElse(0L); + for (Tool tool : config.get().tools().orElse(ImmutableList.of())) { + chars += tool.toJson().length(); + } + } + for (Content content : leading(request.contents(), contentsCount)) { + chars += textLength(content); + } + return chars / CHARS_PER_TOKEN; + } + + private static long textLength(Content content) { + return content.parts().orElse(ImmutableList.of()).stream() + .mapToLong(part -> part.text().map(String::length).orElse(0)) + .sum(); + } + + private Completable deleteCache(String cacheName) { + return Completable.defer( + () -> + Completable.fromFuture( + apiClient.async.caches.delete(cacheName, /* config= */ null))) + .doOnComplete(() -> logger.info("Deleted cache {}.", cacheName)) + .onErrorComplete( + error -> { + if (error instanceof InterruptedException) { + Thread.currentThread().interrupt(); + return false; + } + logger.warn("Failed to delete cache {}: {}", cacheName, describe(error)); + return true; + }); + } + + /** + * Rewrites the request to reference the cache instead of sending what it holds. The final content + * is always sent, because the API rejects a request with no contents. + */ + private static LlmRequest applyCache(LlmRequest request, String cacheName, int contentsCount) { + GenerateContentConfig config = + request.config().orElseGet(() -> GenerateContentConfig.builder().build()).toBuilder() + .clearSystemInstruction() + .clearTools() + .clearToolConfig() + .cachedContent(cacheName) + .build(); + List contents = request.contents(); + int removable = Math.min(contentsCount, Math.max(contents.size() - 1, 0)); + return request.toBuilder() + .config(config) + .contents(ImmutableList.copyOf(contents.subList(removable, contents.size()))) + .build(); + } + + private static List leading(List contents, int count) { + return contents.subList(0, Math.min(count, contents.size())); + } + + /** Describes a failure without its message, which can echo request content. */ + private static String describe(Throwable error) { + Throwable cause = + error instanceof ExecutionException && error.getCause() != null ? error.getCause() : error; + if (cause instanceof ApiException apiException) { + return cause.getClass().getSimpleName() + + " (code " + + apiException.code() + + ", status " + + apiException.status() + + ")"; + } + return cause.getClass().getSimpleName(); + } +} diff --git a/core/src/main/java/com/google/adk/models/LlmRequest.java b/core/src/main/java/com/google/adk/models/LlmRequest.java index d6b56acdb..07c06abc6 100644 --- a/core/src/main/java/com/google/adk/models/LlmRequest.java +++ b/core/src/main/java/com/google/adk/models/LlmRequest.java @@ -25,6 +25,7 @@ import com.fasterxml.jackson.annotation.JsonProperty; import com.fasterxml.jackson.databind.annotation.JsonDeserialize; import com.google.adk.JsonBaseModel; +import com.google.adk.agents.ContextCacheConfig; import com.google.adk.agents.Role; import com.google.adk.tools.BaseTool; import com.google.auto.value.AutoValue; @@ -40,6 +41,7 @@ import java.util.Map; import java.util.Optional; import java.util.stream.Stream; +import org.jspecify.annotations.Nullable; /** Represents a request to be sent to the LLM. */ @AutoValue @@ -88,6 +90,18 @@ public abstract class LlmRequest extends JsonBaseModel { @JsonIgnore public abstract Map tools(); + /** Context cache configuration for this request; empty when context caching is disabled. */ + @JsonIgnore + public abstract Optional cacheConfig(); + + /** Cache state from the agent's latest response, used to reuse or refresh the cache. */ + @JsonIgnore + public abstract Optional cacheMetadata(); + + /** Prompt token count of the agent's previous request, used to decide whether to cache. */ + @JsonIgnore + public abstract Optional cacheableContentsTokenCount(); + /** returns the first system instruction text from the request if present. */ @JsonIgnore public Optional getFirstSystemInstruction() { @@ -155,6 +169,16 @@ private static Builder create() { abstract Map tools(); + @CanIgnoreReturnValue + public abstract Builder cacheConfig(@Nullable ContextCacheConfig cacheConfig); + + @CanIgnoreReturnValue + public abstract Builder cacheMetadata(@Nullable CacheMetadata cacheMetadata); + + @CanIgnoreReturnValue + public abstract Builder cacheableContentsTokenCount( + @Nullable Integer cacheableContentsTokenCount); + @CanIgnoreReturnValue public final Builder appendInstructions(List instructions) { if (instructions.isEmpty()) { diff --git a/core/src/main/java/com/google/adk/models/LlmResponse.java b/core/src/main/java/com/google/adk/models/LlmResponse.java index 5e97eafc2..424e9e570 100644 --- a/core/src/main/java/com/google/adk/models/LlmResponse.java +++ b/core/src/main/java/com/google/adk/models/LlmResponse.java @@ -130,6 +130,10 @@ public abstract class LlmResponse extends JsonBaseModel { @JsonProperty("outputTranscription") public abstract Optional outputTranscription(); + /** Context cache state for this response; set only when context caching is enabled. */ + @JsonProperty("cacheMetadata") + public abstract Optional cacheMetadata(); + public abstract Builder toBuilder(); /** Builder for constructing {@link LlmResponse} instances. */ @@ -185,6 +189,9 @@ public abstract Builder usageMetadata( @JsonProperty("outputTranscription") public abstract Builder outputTranscription(@Nullable Transcription outputTranscription); + @JsonProperty("cacheMetadata") + public abstract Builder cacheMetadata(@Nullable CacheMetadata cacheMetadata); + @CanIgnoreReturnValue public final Builder response(GenerateContentResponse response) { Optional> candidatesOpt = response.candidates(); diff --git a/core/src/test/java/com/google/adk/agents/ContextCacheConfigTest.java b/core/src/test/java/com/google/adk/agents/ContextCacheConfigTest.java new file mode 100644 index 000000000..fcd312aef --- /dev/null +++ b/core/src/test/java/com/google/adk/agents/ContextCacheConfigTest.java @@ -0,0 +1,148 @@ +/* + * 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.agents; + +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertThrows; + +import com.google.common.collect.ImmutableMap; +import com.google.genai.types.HttpOptions; +import java.time.Duration; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public final class ContextCacheConfigTest { + + private static final Duration TTL = Duration.ofMinutes(30); + + @Test + public void defaultConstructor_usesDefaults() { + ContextCacheConfig config = new ContextCacheConfig(); + + assertThat(config.cacheIntervals()).isEqualTo(10); + assertThat(config.ttl()).isEqualTo(Duration.ofMinutes(30)); + assertThat(config.minTokens()).isEqualTo(0); + assertThat(config.createHttpOptions()).isNull(); + } + + @Test + public void threeArgConstructor_leavesHttpOptionsUnset() { + ContextCacheConfig config = new ContextCacheConfig(5, TTL, 100); + + assertThat(config.cacheIntervals()).isEqualTo(5); + assertThat(config.createHttpOptions()).isNull(); + } + + @Test + public void constructor_keepsHttpOptions() { + HttpOptions httpOptions = HttpOptions.builder().timeout(10_000).build(); + + assertThat(new ContextCacheConfig(5, TTL, 100, httpOptions).createHttpOptions()) + .isEqualTo(httpOptions); + } + + @Test + @SuppressWarnings({"deprecation", "InlineMeInliner"}) // Covers the deprecated alias. + public void maxInvocations_returnsCacheIntervals() { + assertThat(new ContextCacheConfig(7, TTL, 0).maxInvocations()).isEqualTo(7); + } + + @Test + public void constructor_cacheIntervalsAtBounds_isAccepted() { + assertThat(new ContextCacheConfig(1, TTL, 0).cacheIntervals()).isEqualTo(1); + assertThat(new ContextCacheConfig(100, TTL, 0).cacheIntervals()).isEqualTo(100); + } + + @Test + public void constructor_cacheIntervalsBelowOne_throws() { + assertThrows(IllegalArgumentException.class, () -> new ContextCacheConfig(0, TTL, 0)); + } + + @Test + public void constructor_cacheIntervalsAboveHundred_throws() { + assertThrows(IllegalArgumentException.class, () -> new ContextCacheConfig(101, TTL, 0)); + } + + @Test + public void constructor_nullTtl_throws() { + assertThrows(NullPointerException.class, () -> new ContextCacheConfig(10, null, 0)); + } + + @Test + public void constructor_zeroTtl_throws() { + assertThrows( + IllegalArgumentException.class, () -> new ContextCacheConfig(10, Duration.ZERO, 0)); + } + + @Test + public void constructor_negativeTtl_throws() { + assertThrows( + IllegalArgumentException.class, + () -> new ContextCacheConfig(10, Duration.ofSeconds(-1), 0)); + } + + @Test + public void constructor_oneSecondTtl_isAccepted() { + assertThat(new ContextCacheConfig(10, Duration.ofSeconds(1), 0).ttl()) + .isEqualTo(Duration.ofSeconds(1)); + } + + @Test + public void constructor_ttlUnderOneSecond_throws() { + assertThrows( + IllegalArgumentException.class, + () -> new ContextCacheConfig(10, Duration.ofMillis(500), 0)); + } + + @Test + public void constructor_fractionalTtl_throws() { + assertThrows( + IllegalArgumentException.class, + () -> new ContextCacheConfig(10, Duration.ofMillis(1500), 0)); + } + + @Test + public void constructor_negativeMinTokens_throws() { + assertThrows(IllegalArgumentException.class, () -> new ContextCacheConfig(10, TTL, -1)); + } + + @Test + public void getTtlString_returnsSeconds() { + assertThat(new ContextCacheConfig(10, Duration.ofMinutes(5), 0).getTtlString()) + .isEqualTo("300s"); + } + + @Test + public void toString_withoutHttpOptions_saysNull() { + assertThat(new ContextCacheConfig().toString()) + .isEqualTo( + "ContextCacheConfig(cacheIntervals=10, ttl=1800s, minTokens=0," + + " createHttpOptions=null)"); + } + + @Test + public void toString_hidesHttpOptionValues() { + HttpOptions httpOptions = + HttpOptions.builder().headers(ImmutableMap.of("x-goog-api-key", "secret")).build(); + + assertThat(new ContextCacheConfig(5, TTL, 100, httpOptions).toString()) + .isEqualTo( + "ContextCacheConfig(cacheIntervals=5, ttl=1800s, minTokens=100," + + " createHttpOptions=set)"); + } +} diff --git a/core/src/test/java/com/google/adk/events/EventTest.java b/core/src/test/java/com/google/adk/events/EventTest.java index 7b91da2c7..7eb3f82e0 100644 --- a/core/src/test/java/com/google/adk/events/EventTest.java +++ b/core/src/test/java/com/google/adk/events/EventTest.java @@ -18,6 +18,7 @@ import static com.google.common.truth.Truth.assertThat; +import com.google.adk.models.CacheMetadata; import com.google.adk.workflow.NodeInfo; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; @@ -349,6 +350,32 @@ public void event_json_serialization_with_transcriptions_works() throws Exceptio assertThat(deserialized.outputTranscription()).hasValue(outputTranscription); } + @Test + public void event_cacheMetadata_survivesJsonAndToBuilder() { + CacheMetadata cacheMetadata = + CacheMetadata.builder() + .fingerprint("abc123") + .contentsCount(2) + .cacheName("cachedContents/42") + .expireTime(Instant.ofEpochSecond(2_000_000_000L, 500_000_000)) + .invocationsUsed(3) + .createdAt(Instant.ofEpochSecond(1_999_998_200L)) + .build(); + Event event = EVENT.toBuilder().cacheMetadata(cacheMetadata).build(); + + assertThat(event.cacheMetadata()).hasValue(cacheMetadata); + assertThat(event.toBuilder().build()).isEqualTo(event); + assertThat(Event.fromJson(event.toJson()).cacheMetadata()).hasValue(cacheMetadata); + assertThat(event).isNotEqualTo(EVENT); + assertThat(event.hashCode()).isNotEqualTo(EVENT.hashCode()); + } + + @Test + public void event_cacheMetadata_emptyByDefault() { + assertThat(EVENT.cacheMetadata()).isEmpty(); + assertThat(EVENT.toJson()).doesNotContain("cacheMetadata"); + } + @Test public void finalResponse_returnsTrueIfNoToolCalls() { Event event = 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 7da99cef0..91bbb8ab1 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 @@ -35,6 +35,7 @@ import com.google.adk.events.Event; import com.google.adk.flows.llmflows.RequestProcessor.RequestProcessingResult; import com.google.adk.flows.llmflows.ResponseProcessor.ResponseProcessingResult; +import com.google.adk.models.CacheMetadata; import com.google.adk.models.LlmRequest; import com.google.adk.models.LlmResponse; import com.google.adk.testing.TestLlm; @@ -59,6 +60,7 @@ import io.reactivex.rxjava3.core.Maybe; import io.reactivex.rxjava3.core.Single; import io.reactivex.rxjava3.schedulers.Schedulers; +import java.time.Instant; import java.util.List; import java.util.Map; import java.util.Optional; @@ -122,6 +124,31 @@ public void run_singleTextResponse_withMetadata_returnsSingleEventWithMetadata() .build()); } + @Test + public void run_responseWithCacheMetadata_copiesItToEvent() { + CacheMetadata cacheMetadata = + CacheMetadata.builder() + .fingerprint("abc123") + .contentsCount(2) + .cacheName("cachedContents/42") + .expireTime(Instant.ofEpochSecond(2_000_000_000L)) + .invocationsUsed(1) + .createdAt(Instant.ofEpochSecond(1_999_998_200L)) + .build(); + TestLlm testLlm = + createTestLlm( + LlmResponse.builder() + .content(Content.fromParts(Part.fromText("LLM response"))) + .cacheMetadata(cacheMetadata) + .build()); + InvocationContext invocationContext = createInvocationContext(createTestAgent(testLlm)); + BaseLlmFlow baseLlmFlow = createBaseLlmFlowWithoutProcessors(); + + List events = baseLlmFlow.run(invocationContext).toList().blockingGet(); + + assertThat(getOnlyElement(events).cacheMetadata()).hasValue(cacheMetadata); + } + @Test public void run_withFunctionCall_returnsCorrectEvents() { Content firstContent = diff --git a/core/src/test/java/com/google/adk/flows/llmflows/ContextCacheRequestProcessorTest.java b/core/src/test/java/com/google/adk/flows/llmflows/ContextCacheRequestProcessorTest.java new file mode 100644 index 000000000..60c233a32 --- /dev/null +++ b/core/src/test/java/com/google/adk/flows/llmflows/ContextCacheRequestProcessorTest.java @@ -0,0 +1,206 @@ +/* + * 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.flows.llmflows; + +import static com.google.adk.testing.TestUtils.createInvocationContext; +import static com.google.adk.testing.TestUtils.createTestAgent; +import static com.google.adk.testing.TestUtils.createTestLlm; +import static com.google.common.truth.Truth.assertThat; + +import com.google.adk.agents.ContextCacheConfig; +import com.google.adk.agents.InvocationContext; +import com.google.adk.events.Event; +import com.google.adk.flows.llmflows.RequestProcessor.RequestProcessingResult; +import com.google.adk.models.CacheMetadata; +import com.google.adk.models.LlmRequest; +import com.google.adk.models.LlmResponse; +import com.google.genai.types.GenerateContentResponseUsageMetadata; +import java.time.Duration; +import java.time.Instant; +import org.jspecify.annotations.Nullable; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public final class ContextCacheRequestProcessorTest { + + private static final ContextCacheConfig CACHE_CONFIG = + new ContextCacheConfig(5, Duration.ofMinutes(10), 1024); + private static final CacheMetadata ACTIVE_CACHE = + CacheMetadata.builder() + .fingerprint("active") + .contentsCount(3) + .cacheName("cachedContents/42") + .expireTime(Instant.ofEpochSecond(2_000_000_000L)) + .invocationsUsed(2) + .createdAt(Instant.ofEpochSecond(1_999_999_000L)) + .build(); + private static final CacheMetadata FINGERPRINT_ONLY = + CacheMetadata.builder().fingerprint("prefix").contentsCount(1).build(); + private static final LlmRequest REQUEST = LlmRequest.builder().model("gemini-2.5-flash").build(); + + private final ContextCacheRequestProcessor processor = new ContextCacheRequestProcessor(); + private final InvocationContext context = + createInvocationContext(createTestAgent(createTestLlm(LlmResponse.builder().build()))) + .toBuilder() + .contextCacheConfig(CACHE_CONFIG) + .build(); + private final String agentName = context.agent().name(); + + @Test + public void processRequest_withoutCacheConfig_returnsRequestUnchanged() { + InvocationContext uncached = + createInvocationContext(createTestAgent(createTestLlm(LlmResponse.builder().build()))); + uncached.session().addEvent(event(agentName, "earlier", ACTIVE_CACHE, 4000)); + + RequestProcessingResult result = processor.processRequest(uncached, REQUEST).blockingGet(); + + assertThat(result.updatedRequest()).isEqualTo(REQUEST); + assertThat(result.events()).isEmpty(); + } + + @Test + public void processRequest_withoutEvents_setsOnlyCacheConfig() { + LlmRequest request = process(); + + assertThat(request.cacheConfig()).hasValue(CACHE_CONFIG); + assertThat(request.cacheMetadata()).isEmpty(); + assertThat(request.cacheableContentsTokenCount()).isEmpty(); + } + + @Test + public void processRequest_activeCacheFromEarlierInvocation_countsOneMoreUse() { + addEvent(event(agentName, "earlier", ACTIVE_CACHE, 4000)); + + LlmRequest request = process(); + + assertThat(request.cacheMetadata()) + .hasValue(ACTIVE_CACHE.toBuilder().invocationsUsed(3).build()); + assertThat(request.cacheableContentsTokenCount()).hasValue(4000); + } + + @Test + public void processRequest_activeCacheFromCurrentInvocation_keepsUseCount() { + addEvent(event(agentName, context.invocationId(), ACTIVE_CACHE, 4000)); + + LlmRequest request = process(); + + assertThat(request.cacheMetadata()).hasValue(ACTIVE_CACHE); + } + + @Test + public void processRequest_eventWithoutInvocationId_keepsUseCount() { + addEvent(event(agentName, /* invocationId= */ null, ACTIVE_CACHE, 4000)); + + LlmRequest request = process(); + + assertThat(request.cacheMetadata()).hasValue(ACTIVE_CACHE); + } + + @Test + public void processRequest_eventWithEmptyInvocationId_keepsUseCount() { + addEvent(event(agentName, "", ACTIVE_CACHE, 4000)); + + LlmRequest request = process(); + + assertThat(request.cacheMetadata()).hasValue(ACTIVE_CACHE); + } + + @Test + public void processRequest_fingerprintOnlyFromEarlierInvocation_isReturnedAsIs() { + addEvent(event(agentName, "earlier", FINGERPRINT_ONLY, 4000)); + + LlmRequest request = process(); + + assertThat(request.cacheMetadata()).hasValue(FINGERPRINT_ONLY); + } + + @Test + public void processRequest_otherAgentsEvents_areIgnored() { + addEvent(event(agentName, "earlier", FINGERPRINT_ONLY, 1000)); + addEvent(event("other agent", "earlier", ACTIVE_CACHE, 9000)); + + LlmRequest request = process(); + + assertThat(request.cacheMetadata()).hasValue(FINGERPRINT_ONLY); + assertThat(request.cacheableContentsTokenCount()).hasValue(1000); + } + + @Test + public void processRequest_severalEvents_usesTheLatest() { + addEvent(event(agentName, "first", FINGERPRINT_ONLY, 1000)); + addEvent(event(agentName, "second", ACTIVE_CACHE, 2000)); + + LlmRequest request = process(); + + assertThat(request.cacheMetadata()) + .hasValue(ACTIVE_CACHE.toBuilder().invocationsUsed(3).build()); + assertThat(request.cacheableContentsTokenCount()).hasValue(2000); + } + + @Test + public void processRequest_latestEventsMissingOneValue_takesItFromOlderEvents() { + addEvent(event(agentName, "first", FINGERPRINT_ONLY, /* promptTokenCount= */ null)); + addEvent(event(agentName, "second", /* cacheMetadata= */ null, 3000)); + + LlmRequest request = process(); + + assertThat(request.cacheMetadata()).hasValue(FINGERPRINT_ONLY); + assertThat(request.cacheableContentsTokenCount()).hasValue(3000); + } + + @Test + public void processRequest_latestEventWithoutUsage_keepsItsMetadataOverOlderMetadata() { + addEvent(event(agentName, "first", ACTIVE_CACHE, 1000)); + addEvent(event(agentName, "second", FINGERPRINT_ONLY, /* promptTokenCount= */ null)); + + LlmRequest request = process(); + + assertThat(request.cacheMetadata()).hasValue(FINGERPRINT_ONLY); + assertThat(request.cacheableContentsTokenCount()).hasValue(1000); + } + + private LlmRequest process() { + RequestProcessingResult result = processor.processRequest(context, REQUEST).blockingGet(); + assertThat(result.events()).isEmpty(); + return result.updatedRequest(); + } + + private void addEvent(Event event) { + context.session().addEvent(event); + } + + private static Event event( + String author, + @Nullable String invocationId, + @Nullable CacheMetadata cacheMetadata, + @Nullable Integer promptTokenCount) { + Event.Builder event = + Event.builder().id(Event.generateEventId()).author(author).cacheMetadata(cacheMetadata); + if (invocationId != null) { + event.invocationId(invocationId); + } + if (promptTokenCount != null) { + event.usageMetadata( + GenerateContentResponseUsageMetadata.builder() + .promptTokenCount(promptTokenCount) + .build()); + } + return event.build(); + } +} diff --git a/core/src/test/java/com/google/adk/flows/llmflows/SingleFlowTest.java b/core/src/test/java/com/google/adk/flows/llmflows/SingleFlowTest.java index ccb10a3a7..e26ffe366 100644 --- a/core/src/test/java/com/google/adk/flows/llmflows/SingleFlowTest.java +++ b/core/src/test/java/com/google/adk/flows/llmflows/SingleFlowTest.java @@ -16,8 +16,16 @@ package com.google.adk.flows.llmflows; +import static com.google.adk.testing.TestUtils.createInvocationContext; +import static com.google.adk.testing.TestUtils.createTestAgent; +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 com.google.adk.agents.ContextCacheConfig; +import com.google.adk.agents.InvocationContext; +import com.google.adk.testing.TestLlm; +import java.time.Duration; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -32,4 +40,25 @@ public void requestProcessors_containsCompaction() { .anyMatch(processor -> processor instanceof Compaction); assertThat(hasCompaction).isTrue(); } + + @Test + public void requestProcessors_runContextCacheAfterContents() { + assertThat(SingleFlow.REQUEST_PROCESSORS.stream().map(Object::getClass)) + .containsAtLeast(Contents.class, ContextCacheRequestProcessor.class) + .inOrder(); + } + + @Test + public void run_withContextCacheConfig_passesItToTheModel() { + ContextCacheConfig cacheConfig = new ContextCacheConfig(5, Duration.ofMinutes(10), 1024); + TestLlm testLlm = createTestLlm(createTextLlmResponse("hi")); + InvocationContext context = + createInvocationContext(createTestAgent(testLlm)).toBuilder() + .contextCacheConfig(cacheConfig) + .build(); + + var unused = new SingleFlow().run(context).toList().blockingGet(); + + assertThat(testLlm.getLastRequest().cacheConfig()).hasValue(cacheConfig); + } } diff --git a/core/src/test/java/com/google/adk/models/CacheMetadataTest.java b/core/src/test/java/com/google/adk/models/CacheMetadataTest.java new file mode 100644 index 000000000..4a99c973d --- /dev/null +++ b/core/src/test/java/com/google/adk/models/CacheMetadataTest.java @@ -0,0 +1,217 @@ +/* + * 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.models; + +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertThrows; + +import com.google.adk.JsonBaseModel; +import com.google.common.collect.ImmutableList; +import java.time.Duration; +import java.time.Instant; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public final class CacheMetadataTest { + + private static final CacheMetadata ACTIVE = + CacheMetadata.builder() + .fingerprint("abc123") + .contentsCount(4) + .cacheName("projects/1/locations/us-central1/cachedContents/42") + .expireTime(Instant.ofEpochSecond(2_000_000_000L, 250_000_000)) + .invocationsUsed(3) + .createdAt(Instant.ofEpochSecond(1_999_998_200L)) + .build(); + + @Test + public void build_fingerprintOnly_hasNoCacheFields() { + CacheMetadata metadata = CacheMetadata.builder().fingerprint("abc123").contentsCount(0).build(); + + assertThat(metadata.fingerprint()).isEqualTo("abc123"); + assertThat(metadata.contentsCount()).isEqualTo(0); + assertThat(metadata.cacheName()).isEmpty(); + assertThat(metadata.expireTime()).isEmpty(); + assertThat(metadata.invocationsUsed()).isEmpty(); + assertThat(metadata.createdAt()).isEmpty(); + } + + @Test + public void build_negativeContentsCount_throws() { + CacheMetadata.Builder builder = CacheMetadata.builder().fingerprint("abc123").contentsCount(-1); + + assertThrows(IllegalStateException.class, builder::build); + } + + @Test + public void build_negativeInvocationsUsed_throws() { + CacheMetadata.Builder builder = ACTIVE.toBuilder().invocationsUsed(-1); + + assertThrows(IllegalStateException.class, builder::build); + } + + @Test + public void build_cacheNameWithoutExpireTime_throws() { + CacheMetadata.Builder builder = ACTIVE.toBuilder().expireTime(null); + + assertThrows(IllegalStateException.class, builder::build); + } + + @Test + public void build_cacheNameWithoutInvocationsUsed_throws() { + CacheMetadata.Builder builder = ACTIVE.toBuilder().invocationsUsed(null); + + assertThrows(IllegalStateException.class, builder::build); + } + + @Test + public void build_expireTimeWithoutCacheName_throws() { + CacheMetadata.Builder builder = ACTIVE.toBuilder().cacheName(null).invocationsUsed(null); + + assertThrows(IllegalStateException.class, builder::build); + } + + @Test + public void build_withoutContentsCount_throws() { + CacheMetadata.Builder builder = CacheMetadata.builder().fingerprint("abc123"); + + assertThrows(IllegalStateException.class, builder::build); + } + + @Test + public void toJson_activeCache_roundTrips() { + String json = ACTIVE.toJson(); + + assertThat(json).contains("\"contentsCount\":4"); + assertThat(JsonBaseModel.fromJsonString(json, CacheMetadata.class)).isEqualTo(ACTIVE); + } + + @Test + public void toJson_fingerprintOnly_omitsCacheFields() throws Exception { + CacheMetadata metadata = CacheMetadata.builder().fingerprint("abc123").contentsCount(2).build(); + + String json = metadata.toJson(); + + assertThat(ImmutableList.copyOf(JsonBaseModel.getMapper().readTree(json).fieldNames())) + .containsExactly("fingerprint", "contentsCount"); + assertThat(JsonBaseModel.fromJsonString(json, CacheMetadata.class)).isEqualTo(metadata); + } + + @Test + public void fromJson_writtenByAdkPython_readsSnakeCaseAndEpochSeconds() { + String json = + "{\"cache_name\": \"projects/1/locations/us-central1/cachedContents/42\"," + + " \"expire_time\": 2000000000.25, \"fingerprint\": \"abc123\"," + + " \"invocations_used\": 3, \"contents_count\": 4, \"created_at\": 1999998200.0}"; + + assertThat(JsonBaseModel.fromJsonString(json, CacheMetadata.class)).isEqualTo(ACTIVE); + } + + @Test + public void fromJson_writtenByAdkKotlin_readsEpochMilliseconds() { + String json = + "{\"cacheName\": \"projects/1/locations/us-central1/cachedContents/42\"," + + " \"expireTime\": 2000000000250, \"fingerprint\": \"abc123\"," + + " \"invocationsUsed\": 3, \"contentsCount\": 4, \"createdAt\": 1999998200000}"; + + assertThat(JsonBaseModel.fromJsonString(json, CacheMetadata.class)).isEqualTo(ACTIVE); + } + + @Test + public void fromJson_isoTimestamps_reads() { + String json = + "{\"cacheName\": \"projects/1/locations/us-central1/cachedContents/42\"," + + " \"expireTime\": \"2033-05-18T03:33:20.250Z\", \"fingerprint\": \"abc123\"," + + " \"invocationsUsed\": 3, \"contentsCount\": 4," + + " \"createdAt\": \"2033-05-18T03:03:20Z\"}"; + + assertThat(JsonBaseModel.fromJsonString(json, CacheMetadata.class)).isEqualTo(ACTIVE); + } + + @Test + public void fromJson_epochAtSecondsMillisBoundary_switchesUnit() { + CacheMetadata below = fingerprintOnlyWithCreatedAt("99999999999"); + CacheMetadata atBoundary = fingerprintOnlyWithCreatedAt("100000000000"); + + assertThat(below.createdAt()).hasValue(Instant.ofEpochSecond(99_999_999_999L)); + assertThat(atBoundary.createdAt()).hasValue(Instant.ofEpochSecond(100_000_000L)); + } + + @Test + public void fromJson_negativeFractionalSeconds_reads() { + assertThat(fingerprintOnlyWithCreatedAt("-1.5").createdAt()) + .hasValue(Instant.ofEpochSecond(-2, 500_000_000)); + } + + @Test + public void fromJson_hugeExponent_throwsWithoutExpandingIt() { + assertThrows(IllegalStateException.class, () -> fingerprintOnlyWithCreatedAt("1e20000000")); + } + + @Test + public void fromJson_negativeEpochMilliseconds_readsByMagnitude() { + assertThat(fingerprintOnlyWithCreatedAt("-1e12").createdAt()) + .hasValue(Instant.ofEpochSecond(-1_000_000_000)); + } + + @Test + public void fromJson_beforeYearOne_throws() { + // -7e13 ms is -7e10 s, before year one. + assertThrows(IllegalStateException.class, () -> fingerprintOnlyWithCreatedAt("-7e13")); + } + + @Test + public void fromJson_malformedTimestampText_throws() { + assertThrows(IllegalStateException.class, () -> fingerprintOnlyWithCreatedAt("\"not-a-time\"")); + } + + @Test + public void fromJson_booleanTimestamp_throws() { + String json = + "{\"cacheName\": \"cachedContents/42\", \"expireTime\": true," + + " \"fingerprint\": \"abc123\", \"invocationsUsed\": 3, \"contentsCount\": 4}"; + + assertThrows( + IllegalStateException.class, () -> JsonBaseModel.fromJsonString(json, CacheMetadata.class)); + } + + private static CacheMetadata fingerprintOnlyWithCreatedAt(String createdAtJson) { + return JsonBaseModel.fromJsonString( + "{\"fingerprint\": \"abc123\", \"contentsCount\": 0, \"createdAt\": " + createdAtJson + "}", + CacheMetadata.class); + } + + @Test + public void toString_fingerprintOnly_showsShortFingerprint() { + CacheMetadata metadata = + CacheMetadata.builder().fingerprint("abc123456789abcd").contentsCount(2).build(); + + assertThat(metadata.toString()) + .isEqualTo("Fingerprint-only: 2 contents, fingerprint=abc12345..."); + } + + @Test + public void toString_activeCache_showsCacheIdAndUse() { + CacheMetadata metadata = + ACTIVE.toBuilder().expireTime(Instant.now().plus(Duration.ofMinutes(10))).build(); + + assertThat(metadata.toString()) + .matches("Cache 42: used 3 invocations, cached 4 contents, expires in (9\\.9|10\\.0)min"); + } +} diff --git a/core/src/test/java/com/google/adk/models/GeminiContextCacheManagerTest.java b/core/src/test/java/com/google/adk/models/GeminiContextCacheManagerTest.java new file mode 100644 index 000000000..c05771416 --- /dev/null +++ b/core/src/test/java/com/google/adk/models/GeminiContextCacheManagerTest.java @@ -0,0 +1,1163 @@ +/* + * 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.models; + +import static com.google.common.collect.ImmutableList.toImmutableList; +import static com.google.common.collect.Iterables.getLast; +import static com.google.common.collect.Streams.stream; +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertThrows; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.MissingNode; +import com.google.adk.agents.ContextCacheConfig; +import com.google.adk.agents.LlmAgent; +import com.google.adk.agents.RunConfig; +import com.google.adk.apps.App; +import com.google.adk.artifacts.InMemoryArtifactService; +import com.google.adk.events.Event; +import com.google.adk.models.GeminiContextCacheManager.CacheResult; +import com.google.adk.runner.Runner; +import com.google.adk.sessions.InMemorySessionService; +import com.google.auth.oauth2.AccessToken; +import com.google.auth.oauth2.GoogleCredentials; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; +import com.google.errorprone.annotations.CanIgnoreReturnValue; +import com.google.genai.Client; +import com.google.genai.types.Candidate; +import com.google.genai.types.ClientOptions; +import com.google.genai.types.Content; +import com.google.genai.types.FinishReason; +import com.google.genai.types.FunctionCallingConfig; +import com.google.genai.types.FunctionCallingConfigMode; +import com.google.genai.types.FunctionDeclaration; +import com.google.genai.types.GenerateContentConfig; +import com.google.genai.types.GenerateContentResponse; +import com.google.genai.types.GenerateContentResponseUsageMetadata; +import com.google.genai.types.GoogleSearch; +import com.google.genai.types.HttpOptions; +import com.google.genai.types.Part; +import com.google.genai.types.Tool; +import com.google.genai.types.ToolConfig; +import java.io.IOException; +import java.time.Duration; +import java.time.Instant; +import java.time.InstantSource; +import java.util.ArrayDeque; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Deque; +import java.util.List; +import okhttp3.Headers; +import okhttp3.Interceptor; +import okhttp3.MediaType; +import okhttp3.OkHttpClient; +import okhttp3.Protocol; +import okhttp3.Request; +import okhttp3.Response; +import okhttp3.ResponseBody; +import okio.Buffer; +import org.jspecify.annotations.Nullable; +import org.junit.After; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Tests for {@link GeminiContextCacheManager}, alone and through {@link Gemini} and a runner. */ +@RunWith(JUnit4.class) +public final class GeminiContextCacheManagerTest { + + private static final ObjectMapper objectMapper = new ObjectMapper(); + private static final String MODEL = "gemini-test-model"; + private static final String INSTRUCTION = "Answer every question briefly and politely."; + private static final Content USER_1 = text("user", "What is the capital of France?"); + private static final Content MODEL_1 = text("model", "Paris."); + private static final Content USER_2 = text("user", "And of Spain?"); + private static final ContextCacheConfig CACHE_CONFIG = + new ContextCacheConfig(10, Duration.ofMinutes(30), 0); + private static final GenerateContentConfig CONFIG = + GenerateContentConfig.builder() + .systemInstruction(Content.fromParts(Part.fromText(INSTRUCTION))) + .tools(functionTool("get_weather", "get_time")) + .toolConfig( + ToolConfig.builder() + .functionCallingConfig( + FunctionCallingConfig.builder().mode(FunctionCallingConfigMode.Known.AUTO)) + .build()) + .temperature(0.2f) + .build(); + private static final String EXISTING_CACHE = "cachedContents/existing"; + private static final Tool SEARCH_TOOL = + Tool.builder().googleSearch(GoogleSearch.builder().build()).build(); + private static final String USER_ID = "user"; + private static final Instant NOW = Instant.parse("2026-10-07T12:00:00Z"); + + private final FakeGeminiApi api = new FakeGeminiApi(); + private final GeminiContextCacheManager manager = + new GeminiContextCacheManager(api.client(), MODEL, InstantSource.fixed(NOW)); + + @After + public void clearInterruptFlag() { + // The interrupt tests set the flag; a failed assertion must not leak it into later tests. + boolean unused = Thread.interrupted(); + } + + @Test + public void handleContextCaching_withoutCacheConfig_fails() { + LlmRequest request = + LlmRequest.builder().model(MODEL).contents(ImmutableList.of(USER_1)).build(); + + assertThrows(IllegalArgumentException.class, () -> handle(request)); + } + + @Test + public void handleContextCaching_withoutPreviousMetadata_onlyFingerprintsThePrefix() { + LlmRequest request = request(USER_1, MODEL_1, USER_2).build(); + + CacheResult result = handle(request); + + assertThat(result.request()).isEqualTo(request); + assertThat(result.metadata().contentsCount()).isEqualTo(2); + assertThat(result.metadata().fingerprint()).matches("[0-9a-f]{16}"); + assertThat(result.metadata().cacheName()).isEmpty(); + assertThat(api.kinds()).isEmpty(); + } + + @Test + public void handleContextCaching_upperCaseUserRole_isNotCached() { + LlmRequest request = request(USER_1, MODEL_1, text("USER", "And of Spain?")).build(); + + assertThat(handle(request).metadata().contentsCount()).isEqualTo(2); + } + + @Test + public void handleContextCaching_validCache_reusesItWithoutCallingTheApi() { + LlmRequest request = request(USER_1, MODEL_1, USER_2).build(); + CacheMetadata cache = activeCacheFor(request, inOneHour(), /* invocationsUsed= */ 3); + + CacheResult result = handle(request.toBuilder().cacheMetadata(cache).build()); + + assertThat(result.metadata()).isEqualTo(cache); + GenerateContentConfig config = result.request().config().get(); + assertThat(config.cachedContent()).hasValue(EXISTING_CACHE); + assertThat(config.systemInstruction()).isEmpty(); + assertThat(config.tools()).isEmpty(); + assertThat(config.toolConfig()).isEmpty(); + assertThat(config.temperature()).hasValue(0.2f); + assertThat(result.request().contents()).containsExactly(USER_2); + assertThat(api.kinds()).isEmpty(); + } + + @Test + public void handleContextCaching_cacheUsedForCacheIntervals_isStillReused() { + LlmRequest request = request(USER_1, MODEL_1, USER_2).build(); + CacheMetadata cache = activeCacheFor(request, inOneHour(), CACHE_CONFIG.cacheIntervals()); + + CacheResult result = handle(request.toBuilder().cacheMetadata(cache).build()); + + assertThat(result.metadata()).isEqualTo(cache); + assertThat(api.kinds()).isEmpty(); + } + + @Test + public void handleContextCaching_cacheUsedForMoreThanCacheIntervals_isReplaced() { + LlmRequest request = request(USER_1, MODEL_1, USER_2).cacheableContentsTokenCount(5000).build(); + CacheMetadata cache = activeCacheFor(request, inOneHour(), CACHE_CONFIG.cacheIntervals() + 1); + + CacheResult result = handle(request.toBuilder().cacheMetadata(cache).build()); + + assertThat(api.kinds()).containsExactly(Kind.DELETE_CACHE, Kind.CREATE_CACHE).inOrder(); + assertThat(result.metadata().cacheName()).hasValue("cachedContents/cache-1"); + assertThat(result.metadata().invocationsUsed()).hasValue(1); + } + + @Test + public void handleContextCaching_expiredCache_isDeletedAndReplaced() { + LlmRequest request = request(USER_1, MODEL_1, USER_2).cacheableContentsTokenCount(5000).build(); + CacheMetadata cache = activeCacheFor(request, NOW.minusSeconds(1), /* invocationsUsed= */ 2); + + CacheResult result = handle(request.toBuilder().cacheMetadata(cache).build()); + + assertThat(api.kinds()).containsExactly(Kind.DELETE_CACHE, Kind.CREATE_CACHE).inOrder(); + assertThat(api.paths(Kind.DELETE_CACHE).get(0)).endsWith("/" + EXISTING_CACHE); + assertThat(result.metadata().cacheName()).hasValue("cachedContents/cache-1"); + assertThat(result.metadata().contentsCount()).isEqualTo(2); + assertThat(result.metadata().invocationsUsed()).hasValue(1); + assertThat(result.request().config().get().cachedContent()).hasValue("cachedContents/cache-1"); + assertThat(result.request().contents()).containsExactly(USER_2); + } + + @Test + public void handleContextCaching_changedInstruction_deletesCacheAndStartsNewFingerprint() { + LlmRequest original = + request(USER_1, MODEL_1, USER_2).cacheableContentsTokenCount(5000).build(); + CacheMetadata cache = activeCacheFor(original, inOneHour(), /* invocationsUsed= */ 2); + LlmRequest changed = + original.toBuilder() + .config( + CONFIG.toBuilder() + .systemInstruction(Content.fromParts(Part.fromText("Answer in French."))) + .build()) + .cacheMetadata(cache) + .build(); + + CacheResult result = handle(changed); + + assertThat(api.kinds()).containsExactly(Kind.DELETE_CACHE); + assertThat(result.request()).isEqualTo(changed); + assertThat(result.metadata().cacheName()).isEmpty(); + assertThat(result.metadata().contentsCount()).isEqualTo(2); + assertThat(result.metadata().fingerprint()).isNotEqualTo(cache.fingerprint()); + } + + @Test + public void handleContextCaching_createHttpOptions_areSentWithTheCreateCall() { + ContextCacheConfig withHeader = + new ContextCacheConfig( + 10, + Duration.ofMinutes(30), + 0, + HttpOptions.builder().headers(ImmutableMap.of("x-adk-test", "create")).build()); + CacheMetadata firstTurn = handle(request(USER_1).cacheConfig(withHeader).build()).metadata(); + + var unused = + handle( + request(USER_1, MODEL_1, USER_2) + .cacheConfig(withHeader) + .cacheMetadata(firstTurn) + .cacheableContentsTokenCount(5000) + .build()); + + assertThat(api.headers(Kind.CREATE_CACHE).get(0).get("x-adk-test")).isEqualTo("create"); + } + + @Test + public void handleContextCaching_unchangedPrefix_cachesTheGrownPrefix() throws Exception { + CacheMetadata firstTurn = handle(request(USER_1).build()).metadata(); + LlmRequest request = + request(USER_1, MODEL_1, USER_2) + .cacheMetadata(firstTurn) + .cacheableContentsTokenCount(5000) + .build(); + + CacheResult result = handle(request); + + assertThat(api.kinds()).containsExactly(Kind.CREATE_CACHE); + JsonNode create = api.bodies(Kind.CREATE_CACHE).get(0); + assertThat(create.get("model").asText()).endsWith(MODEL); + assertThat(texts(create.get("contents"))) + .containsExactlyElementsIn(texts(USER_1, MODEL_1)) + .inOrder(); + assertThat(create.get("systemInstruction").toString()).contains(INSTRUCTION); + assertThat(create.get("tools").toString()).contains("get_weather"); + assertThat(create.get("toolConfig").toString()).contains("AUTO"); + assertThat(create.get("ttl").asText()).isEqualTo("1800s"); + assertThat(create.get("displayName").asText()) + .isEqualTo("adk-cache-" + NOW.getEpochSecond() + "-2contents"); + CacheMetadata created = result.metadata(); + assertThat(created.cacheName()).hasValue("cachedContents/cache-1"); + assertThat(created.contentsCount()).isEqualTo(2); + assertThat(created.invocationsUsed()).hasValue(1); + assertThat(created.expireTime()).hasValue(api.expireTime(1)); + assertThat(result.request().contents()).containsExactly(USER_2); + // The created fingerprint validates the same request, so the next invocation reuses it. + assertThat(handle(request.toBuilder().cacheMetadata(created).build()).metadata()) + .isEqualTo(created); + assertThat(api.kinds()).containsExactly(Kind.CREATE_CACHE); + } + + @Test + public void handleContextCaching_withoutPreviousTokenCount_keepsGrownFingerprint() { + CacheMetadata firstTurn = handle(request(USER_1).build()).metadata(); + LlmRequest request = request(USER_1, MODEL_1, USER_2).cacheMetadata(firstTurn).build(); + + CacheResult result = handle(request); + + assertThat(api.kinds()).isEmpty(); + assertThat(result.request()).isEqualTo(request); + assertThat(result.metadata().cacheName()).isEmpty(); + assertThat(result.metadata().contentsCount()).isEqualTo(2); + assertThat(result.metadata().fingerprint()).isNotEqualTo(firstTurn.fingerprint()); + } + + @Test + public void handleContextCaching_previousRequestBelowMinTokens_doesNotCreateCache() { + CacheMetadata firstTurn = handle(request(USER_1).build()).metadata(); + LlmRequest request = + request(USER_1, MODEL_1, USER_2) + .cacheConfig(new ContextCacheConfig(10, Duration.ofMinutes(30), 6000)) + .cacheMetadata(firstTurn) + .cacheableContentsTokenCount(5000) + .build(); + + CacheResult result = handle(request); + + assertThat(api.kinds()).isEmpty(); + assertThat(result.metadata().cacheName()).isEmpty(); + assertThat(result.metadata().contentsCount()).isEqualTo(2); + } + + @Test + public void handleContextCaching_gemini25PrefixAboveMinimum_createsCache() { + assertThat(createsCache("gemini-2.5-flash", /* previousTokenCount= */ 3000)).isTrue(); + } + + @Test + public void handleContextCaching_gemini25PrefixBelowMinimum_doesNotCreateCache() { + assertThat(createsCache("gemini-2.5-flash", /* previousTokenCount= */ 2000)).isFalse(); + } + + @Test + public void handleContextCaching_gemini3PrefixBelowMinimum_doesNotCreateCache() { + assertThat(createsCache("gemini-3-pro", /* previousTokenCount= */ 4000)).isFalse(); + } + + @Test + public void handleContextCaching_gemini3ResourceNameBelowMinimum_doesNotCreateCache() { + assertThat( + createsCache( + "projects/p/locations/l/models/gemini-3-pro", /* previousTokenCount= */ 4000)) + .isFalse(); + } + + @Test + public void handleContextCaching_opaqueModel_appliesNoMinimum() { + assertThat(createsCache("projects/p/locations/l/endpoints/123", /* previousTokenCount= */ 10)) + .isTrue(); + } + + @Test + public void handleContextCaching_smallPrefixOfLargePrompt_doesNotCreateCache() { + String model = "gemini-2.5-flash"; + GeminiContextCacheManager manager = new GeminiContextCacheManager(api.client(), model); + Content largeQuestion = text("user", "x".repeat(400_000)); + CacheMetadata firstTurn = + manager.handleContextCaching(request(USER_1).model(model).build()).blockingGet().metadata(); + LlmRequest request = + request(USER_1, MODEL_1, largeQuestion) + .model(model) + .cacheMetadata(firstTurn) + .cacheableContentsTokenCount(100_000) + .build(); + + CacheResult result = manager.handleContextCaching(request).blockingGet(); + + assertThat(api.kinds()).isEmpty(); + assertThat(result.metadata().cacheName()).isEmpty(); + } + + @Test + public void handleContextCaching_onlyBinaryParts_trustsPreviousTokenCount() { + String model = "gemini-2.5-flash"; + GeminiContextCacheManager manager = new GeminiContextCacheManager(api.client(), model); + Content image = + Content.builder() + .role("user") + .parts(Part.fromBytes(new byte[] {1, 2}, "image/png")) + .build(); + Content answer = + Content.builder().role("model").parts(Part.fromBytes(new byte[] {3}, "image/png")).build(); + LlmRequest firstTurn = + LlmRequest.builder() + .model(model) + .contents(ImmutableList.of(image)) + .cacheConfig(CACHE_CONFIG) + .build(); + CacheMetadata firstMetadata = manager.handleContextCaching(firstTurn).blockingGet().metadata(); + LlmRequest request = + firstTurn.toBuilder() + .contents(ImmutableList.of(image, answer, image)) + .cacheMetadata(firstMetadata) + .cacheableContentsTokenCount(3000) + .build(); + + CacheResult result = manager.handleContextCaching(request).blockingGet(); + + assertThat(api.kinds()).containsExactly(Kind.CREATE_CACHE); + JsonNode create = api.bodies(Kind.CREATE_CACHE).get(0); + assertThat(create.has("systemInstruction")).isFalse(); + assertThat(create.has("tools")).isFalse(); + assertThat(create.get("contents")).hasSize(2); + assertThat(result.request().config().get().cachedContent()).hasValue("cachedContents/cache-1"); + assertThat(result.request().contents()).containsExactly(image); + } + + @Test + public void handleContextCaching_creationFails_sendsRequestUncachedWithGrownFingerprint() { + api.failNextCreate(400); + CacheMetadata firstTurn = handle(request(USER_1).build()).metadata(); + LlmRequest request = + request(USER_1, MODEL_1, USER_2) + .cacheMetadata(firstTurn) + .cacheableContentsTokenCount(5000) + .build(); + + CacheResult result = handle(request); + + assertThat(api.kinds()).containsExactly(Kind.CREATE_CACHE); + assertThat(result.request()).isEqualTo(request); + assertThat(result.metadata().cacheName()).isEmpty(); + assertThat(result.metadata().contentsCount()).isEqualTo(2); + // The grown prefix is kept, so the next invocation tries to create the cache again. + var unused = handle(request.toBuilder().cacheMetadata(result.metadata()).build()); + assertThat(api.kinds()).containsExactly(Kind.CREATE_CACHE, Kind.CREATE_CACHE); + } + + @Test + public void handleContextCaching_deletionFails_stillCreatesNewCache() { + api.failNextDelete(500); + LlmRequest request = request(USER_1, MODEL_1, USER_2).cacheableContentsTokenCount(5000).build(); + CacheMetadata cache = activeCacheFor(request, NOW.minusSeconds(1), /* invocationsUsed= */ 2); + + CacheResult result = handle(request.toBuilder().cacheMetadata(cache).build()); + + assertThat(api.kinds()).containsExactly(Kind.DELETE_CACHE, Kind.CREATE_CACHE).inOrder(); + assertThat(result.metadata().cacheName()).hasValue("cachedContents/cache-1"); + } + + @Test + public void handleContextCaching_serverReturnsEmptyCacheName_sendsRequestUncached() { + api.answerNextCreateWith("{\"name\": \"\"}"); + CacheMetadata firstTurn = handle(request(USER_1).build()).metadata(); + LlmRequest request = + request(USER_1, MODEL_1, USER_2) + .cacheMetadata(firstTurn) + .cacheableContentsTokenCount(5000) + .build(); + + CacheResult result = handle(request); + + assertThat(api.kinds()).containsExactly(Kind.CREATE_CACHE); + assertThat(result.request()).isEqualTo(request); + assertThat(result.metadata().cacheName()).isEmpty(); + } + + @Test + public void handleContextCaching_cacheExpiringNow_isReplaced() { + LlmRequest request = request(USER_1, MODEL_1, USER_2).cacheableContentsTokenCount(5000).build(); + CacheMetadata cache = activeCacheFor(request, NOW, /* invocationsUsed= */ 1); + + var unused = handle(request.toBuilder().cacheMetadata(cache).build()); + + assertThat(api.kinds()).containsExactly(Kind.DELETE_CACHE, Kind.CREATE_CACHE).inOrder(); + } + + @Test + public void handleContextCaching_previousRequestAtMinTokens_createsCache() { + CacheMetadata firstTurn = handle(request(USER_1).build()).metadata(); + LlmRequest request = + request(USER_1, MODEL_1, USER_2) + .cacheConfig(new ContextCacheConfig(10, Duration.ofMinutes(30), 5000)) + .cacheMetadata(firstTurn) + .cacheableContentsTokenCount(5000) + .build(); + + CacheResult result = handle(request); + + assertThat(api.kinds()).containsExactly(Kind.CREATE_CACHE); + assertThat(result.metadata().cacheName()).hasValue("cachedContents/cache-1"); + } + + @Test + public void handleContextCaching_previousCountAboveCurrent_keepsPreviousCount() { + LlmRequest request = request(USER_1, MODEL_1).build(); + CacheMetadata previous = + CacheMetadata.builder().fingerprint(fingerprint(request)).contentsCount(5).build(); + + CacheResult result = + handle( + request.toBuilder().cacheMetadata(previous).cacheableContentsTokenCount(5000).build()); + + assertThat(api.bodies(Kind.CREATE_CACHE).get(0).get("contents")).hasSize(2); + assertThat(result.metadata().contentsCount()).isEqualTo(5); + assertThat(result.request().contents()).containsExactly(MODEL_1); + } + + @Test + public void handleContextCaching_interruptedWhileCreating_propagatesAndKeepsInterrupt() { + CacheMetadata firstTurn = handle(request(USER_1).build()).metadata(); + LlmRequest request = + request(USER_1, MODEL_1, USER_2) + .cacheMetadata(firstTurn) + .cacheableContentsTokenCount(5000) + .build(); + api.delayCacheCalls(Duration.ofSeconds(2)); + + Thread.currentThread().interrupt(); + RuntimeException thrown = assertThrows(RuntimeException.class, () -> handle(request)); + + assertThat(Thread.interrupted()).isTrue(); + assertThat(thrown).hasCauseThat().isInstanceOf(InterruptedException.class); + } + + @Test + public void handleContextCaching_interruptedWhileDeleting_propagatesAndKeepsInterrupt() { + LlmRequest request = request(USER_1, MODEL_1, USER_2).cacheableContentsTokenCount(5000).build(); + CacheMetadata cache = activeCacheFor(request, NOW.minusSeconds(1), /* invocationsUsed= */ 1); + LlmRequest withExpiredCache = request.toBuilder().cacheMetadata(cache).build(); + api.delayCacheCalls(Duration.ofSeconds(2)); + + Thread.currentThread().interrupt(); + RuntimeException thrown = assertThrows(RuntimeException.class, () -> handle(withExpiredCache)); + + assertThat(Thread.interrupted()).isTrue(); + assertThat(thrown).hasCauseThat().isInstanceOf(InterruptedException.class); + } + + @Test + public void handleContextCaching_serverOmitsExpireTime_expiresAfterTtl() { + api.omitExpireTime(); + CacheMetadata firstTurn = handle(request(USER_1).build()).metadata(); + + CacheMetadata created = + handle( + request(USER_1, MODEL_1, USER_2) + .cacheMetadata(firstTurn) + .cacheableContentsTokenCount(5000) + .build()) + .metadata(); + + assertThat(created.createdAt()).hasValue(NOW); + assertThat(created.expireTime()).hasValue(NOW.plus(CACHE_CONFIG.ttl())); + } + + @Test + public void handleContextCaching_cacheCoversEveryContent_stillSendsTheFinalContent() { + LlmRequest request = request(USER_1, MODEL_1).build(); + CacheMetadata cache = activeCacheFor(request, inOneHour(), /* invocationsUsed= */ 1); + + CacheResult result = handle(request.toBuilder().cacheMetadata(cache).build()); + + assertThat(cache.contentsCount()).isEqualTo(2); + assertThat(result.request().contents()).containsExactly(MODEL_1); + } + + @Test + public void handleContextCaching_requestWithoutConfig_cachesOnlyContents() { + LlmRequest firstTurn = + LlmRequest.builder() + .model(MODEL) + .contents(ImmutableList.of(USER_1)) + .cacheConfig(CACHE_CONFIG) + .build(); + CacheMetadata firstMetadata = handle(firstTurn).metadata(); + + CacheResult result = + handle( + firstTurn.toBuilder() + .contents(ImmutableList.of(USER_1, MODEL_1, USER_2)) + .cacheMetadata(firstMetadata) + .cacheableContentsTokenCount(5000) + .build()); + + JsonNode create = api.bodies(Kind.CREATE_CACHE).get(0); + assertThat(create.has("systemInstruction")).isFalse(); + assertThat(create.has("tools")).isFalse(); + assertThat(create.has("toolConfig")).isFalse(); + assertThat(result.request().config()) + .hasValue(GenerateContentConfig.builder().cachedContent("cachedContents/cache-1").build()); + } + + @Test + public void handleContextCaching_serverReturnsNoCacheName_sendsRequestUncached() { + api.answerNextCreateWith("{}"); + CacheMetadata firstTurn = handle(request(USER_1).build()).metadata(); + LlmRequest request = + request(USER_1, MODEL_1, USER_2) + .cacheMetadata(firstTurn) + .cacheableContentsTokenCount(5000) + .build(); + + CacheResult result = handle(request); + + assertThat(api.kinds()).containsExactly(Kind.CREATE_CACHE); + assertThat(result.request()).isEqualTo(request); + assertThat(result.metadata().cacheName()).isEmpty(); + assertThat(result.metadata().contentsCount()).isEqualTo(2); + } + + @Test + public void handleContextCaching_noSettledContents_cachesOnlyInstructionAndTools() { + CacheMetadata firstTurn = handle(request(USER_1).build()).metadata(); + Content followUp = text("user", "Here is some extra context."); + + CacheResult result = + handle( + request(USER_1, followUp) + .cacheMetadata(firstTurn) + .cacheableContentsTokenCount(5000) + .build()); + + JsonNode create = api.bodies(Kind.CREATE_CACHE).get(0); + assertThat(create.has("contents")).isFalse(); + assertThat(create.get("systemInstruction").toString()).contains(INSTRUCTION); + assertThat(result.metadata().contentsCount()).isEqualTo(0); + assertThat(result.request().contents()).containsExactly(USER_1, followUp).inOrder(); + assertThat(result.request().config().get().systemInstruction()).isEmpty(); + } + + @Test + public void handleContextCaching_emptyToolList_isTreatedAsNoTools() { + GenerateContentConfig emptyTools = CONFIG.toBuilder().tools(ImmutableList.of()).build(); + GenerateContentConfig noTools = CONFIG.toBuilder().clearTools().build(); + CacheMetadata firstTurn = handle(request(USER_1).config(emptyTools).build()).metadata(); + + var unused = + handle( + request(USER_1, MODEL_1, USER_2) + .config(emptyTools) + .cacheMetadata(firstTurn) + .cacheableContentsTokenCount(5000) + .build()); + + assertThat(firstTurn.fingerprint()) + .isEqualTo(fingerprint(request(USER_1).config(noTools).build())); + assertThat(api.bodies(Kind.CREATE_CACHE).get(0).has("tools")).isFalse(); + } + + @Test + public void handleContextCaching_gemini25WithZeroPreviousTokens_doesNotCreateCache() { + assertThat(createsCache("gemini-2.5-flash", /* previousTokenCount= */ 0)).isFalse(); + } + + @Test + public void fingerprint_trailingUserContent_isIgnored() { + assertThat(fingerprint(request(USER_1, MODEL_1, USER_2).build())) + .isEqualTo(fingerprint(request(USER_1, MODEL_1, text("user", "Other question")).build())); + } + + @Test + public void fingerprint_reorderedTools_isUnchanged() { + LlmRequest request = + request(USER_1) + .config( + CONFIG.toBuilder() + .tools(functionTool("b_tool", "a_tool"), SEARCH_TOOL, functionTool("c_tool")) + .build()) + .build(); + LlmRequest reordered = + request(USER_1) + .config( + CONFIG.toBuilder() + .tools(functionTool("c_tool"), functionTool("a_tool", "b_tool"), SEARCH_TOOL) + .build()) + .build(); + + assertThat(fingerprint(reordered)).isEqualTo(fingerprint(request)); + } + + @Test + public void fingerprint_differentTools_differs() { + LlmRequest request = request(USER_1).build(); + LlmRequest otherTools = + request.toBuilder().config(CONFIG.toBuilder().tools(functionTool("other")).build()).build(); + + assertThat(fingerprint(otherTools)).isNotEqualTo(fingerprint(request)); + } + + @Test + public void fingerprint_differentToolConfig_differs() { + LlmRequest request = request(USER_1).build(); + LlmRequest otherToolConfig = + request.toBuilder() + .config( + CONFIG.toBuilder() + .toolConfig( + ToolConfig.builder() + .functionCallingConfig( + FunctionCallingConfig.builder() + .mode(FunctionCallingConfigMode.Known.NONE)) + .build()) + .build()) + .build(); + + assertThat(fingerprint(otherToolConfig)).isNotEqualTo(fingerprint(request)); + } + + @Test + public void fingerprint_differentModel_differs() { + LlmRequest request = request(USER_1).build(); + GeminiContextCacheManager otherModel = + new GeminiContextCacheManager(api.client(), "gemini-other-model"); + + assertThat(otherModel.handleContextCaching(request).blockingGet().metadata().fingerprint()) + .isNotEqualTo(fingerprint(request)); + } + + @Test + public void fingerprint_differentVertexProjectOrLocation_differs() { + LlmRequest request = request(USER_1).build(); + + String base = vertexFingerprint("project-a", "us-central1", request); + + assertThat(vertexFingerprint("project-b", "us-central1", request)).isNotEqualTo(base); + assertThat(vertexFingerprint("project-a", "europe-west4", request)).isNotEqualTo(base); + assertThat(vertexFingerprint("project-a", "us-central1", request)).isEqualTo(base); + } + + @Test + public void fingerprint_differentBackend_differs() { + LlmRequest request = request(USER_1).build(); + Client vertexClient = Client.builder().vertexAI(true).apiKey("test-api-key").build(); + GeminiContextCacheManager vertex = new GeminiContextCacheManager(vertexClient, MODEL); + + assertThat(vertex.handleContextCaching(request).blockingGet().metadata().fingerprint()) + .isNotEqualTo(fingerprint(request)); + } + + @Test + public void runAsync_secondTurn_createsCacheAndSendsOnlyUncachedContents() { + Runner runner = runner(CACHE_CONFIG); + String sessionId = newSession(runner); + + run(runner, sessionId, "Hello"); + run(runner, sessionId, "Tell me more"); + + assertThat(api.kinds()) + .containsExactly(Kind.GENERATE, Kind.CREATE_CACHE, Kind.GENERATE) + .inOrder(); + JsonNode createCache = api.bodies(Kind.CREATE_CACHE).get(0); + assertThat(texts(createCache.get("contents"))).containsExactly("Hello", "Answer 1").inOrder(); + assertThat(createCache.get("systemInstruction").toString()).contains(INSTRUCTION); + assertThat(createCache.get("ttl").asText()).isEqualTo("1800s"); + JsonNode secondCall = api.bodies(Kind.GENERATE).get(1); + assertThat(secondCall.path("cachedContent").asText()).isEqualTo("cachedContents/cache-1"); + assertThat(secondCall.has("systemInstruction")).isFalse(); + assertThat(texts(secondCall.get("contents"))).containsExactly("Tell me more"); + } + + @Test + public void runAsync_thirdTurn_reusesCacheWithoutCreatingAnother() { + Runner runner = runner(CACHE_CONFIG); + String sessionId = newSession(runner); + + run(runner, sessionId, "Hello"); + Event secondAnswer = getLast(run(runner, sessionId, "Tell me more")); + Event thirdAnswer = getLast(run(runner, sessionId, "And then?")); + + assertThat(api.kinds()) + .containsExactly(Kind.GENERATE, Kind.CREATE_CACHE, Kind.GENERATE, Kind.GENERATE) + .inOrder(); + JsonNode thirdCall = api.bodies(Kind.GENERATE).get(2); + assertThat(thirdCall.path("cachedContent").asText()).isEqualTo("cachedContents/cache-1"); + assertThat(thirdCall.has("systemInstruction")).isFalse(); + assertThat(texts(thirdCall.get("contents"))) + .containsExactly("Tell me more", "Answer 2", "And then?") + .inOrder(); + assertThat(secondAnswer.cacheMetadata().get().invocationsUsed()).hasValue(1); + assertThat(thirdAnswer.cacheMetadata().get().cacheName()).hasValue("cachedContents/cache-1"); + assertThat(thirdAnswer.cacheMetadata().get().invocationsUsed()).hasValue(2); + } + + @Test + public void runAsync_streaming_putsCacheMetadataOnlyOnTheFinalEvent() { + Runner runner = runner(CACHE_CONFIG); + String sessionId = newSession(runner); + RunConfig streaming = RunConfig.builder().streamingMode(RunConfig.StreamingMode.SSE).build(); + + run(runner, sessionId, "Hello", streaming); + ImmutableList events = run(runner, sessionId, "Tell me more", streaming); + + assertThat(api.kinds()) + .containsExactly(Kind.GENERATE, Kind.CREATE_CACHE, Kind.GENERATE) + .inOrder(); + assertThat(api.bodies(Kind.GENERATE).get(1).path("cachedContent").asText()) + .isEqualTo("cachedContents/cache-1"); + ImmutableList partials = + events.stream().filter(event -> event.partial().orElse(false)).collect(toImmutableList()); + assertThat(partials).isNotEmpty(); + assertThat( + partials.stream() + .filter(event -> event.cacheMetadata().isPresent()) + .collect(toImmutableList())) + .isEmpty(); + assertThat(getLast(events).partial()).isEmpty(); + assertThat(getLast(events).cacheMetadata().get().cacheName()) + .hasValue("cachedContents/cache-1"); + } + + @Test + public void runAsync_turnThatCreatesCache_staysOnTheCallerThread() { + api.delayCacheCalls(Duration.ofMillis(300)); + Runner runner = runner(CACHE_CONFIG); + String sessionId = newSession(runner); + + ImmutableList firstTurn = + eventThreads(runner, sessionId, "Hello", RunConfig.builder().build()); + ImmutableList secondTurn = + eventThreads(runner, sessionId, "Tell me more", RunConfig.builder().build()); + + assertThat(api.kinds()).contains(Kind.CREATE_CACHE); + assertThat(firstTurn).isNotEmpty(); + assertThat(ImmutableSet.copyOf(secondTurn)) + .containsExactlyElementsIn(ImmutableSet.copyOf(firstTurn)); + } + + @Test + public void runAsync_streamingTurnThatCreatesCache_staysOnTheCallerThread() { + api.delayCacheCalls(Duration.ofMillis(300)); + Runner runner = runner(CACHE_CONFIG); + String sessionId = newSession(runner); + RunConfig streaming = RunConfig.builder().streamingMode(RunConfig.StreamingMode.SSE).build(); + + ImmutableList firstTurn = eventThreads(runner, sessionId, "Hello", streaming); + ImmutableList secondTurn = eventThreads(runner, sessionId, "Tell me more", streaming); + + assertThat(api.kinds()).contains(Kind.CREATE_CACHE); + assertThat(firstTurn).isNotEmpty(); + assertThat(ImmutableSet.copyOf(secondTurn)) + .containsExactlyElementsIn(ImmutableSet.copyOf(firstTurn)); + } + + @Test + public void runAsync_withoutCacheConfig_neverCreatesCache() { + Runner runner = runner(/* cacheConfig= */ null); + String sessionId = newSession(runner); + + run(runner, sessionId, "Hello"); + Event secondAnswer = getLast(run(runner, sessionId, "Tell me more")); + + assertThat(api.kinds()).containsExactly(Kind.GENERATE, Kind.GENERATE); + JsonNode secondCall = api.bodies(Kind.GENERATE).get(1); + assertThat(secondCall.has("cachedContent")).isFalse(); + assertThat(secondCall.get("systemInstruction").toString()).contains(INSTRUCTION); + assertThat(texts(secondCall.get("contents"))) + .containsExactly("Hello", "Answer 1", "Tell me more") + .inOrder(); + assertThat(secondAnswer.cacheMetadata()).isEmpty(); + } + + private static String vertexFingerprint(String project, String location, LlmRequest request) { + Client client = + Client.builder() + .vertexAI(true) + .project(project) + .location(location) + .credentials(GoogleCredentials.create(new AccessToken("test-token", null))) + .build(); + return new GeminiContextCacheManager(client, MODEL) + .handleContextCaching(request) + .blockingGet() + .metadata() + .fingerprint(); + } + + private CacheResult handle(LlmRequest request) { + return manager.handleContextCaching(request).blockingGet(); + } + + private String fingerprint(LlmRequest request) { + return handle(request).metadata().fingerprint(); + } + + /** Returns metadata for an existing cache of the request's cacheable prefix. */ + private CacheMetadata activeCacheFor( + LlmRequest request, Instant expireTime, int invocationsUsed) { + return handle(request.toBuilder().cacheMetadata(null).build()).metadata().toBuilder() + .cacheName(EXISTING_CACHE) + .expireTime(expireTime) + .invocationsUsed(invocationsUsed) + .createdAt(NOW.minus(Duration.ofMinutes(5))) + .build(); + } + + /** Returns whether the second turn of a conversation with {@code model} creates a cache. */ + private boolean createsCache(String model, int previousTokenCount) { + GeminiContextCacheManager manager = new GeminiContextCacheManager(api.client(), model); + CacheMetadata firstTurn = + manager.handleContextCaching(request(USER_1).model(model).build()).blockingGet().metadata(); + CacheResult result = + manager + .handleContextCaching( + request(USER_1, MODEL_1, USER_2) + .model(model) + .cacheMetadata(firstTurn) + .cacheableContentsTokenCount(previousTokenCount) + .build()) + .blockingGet(); + return result.metadata().cacheName().isPresent(); + } + + private Runner runner(@Nullable ContextCacheConfig cacheConfig) { + LlmAgent agent = + LlmAgent.builder() + .name("assistant") + .model(new Gemini(MODEL, api.client())) + .instruction(INSTRUCTION) + .build(); + App.Builder app = App.builder().name("cache_app").rootAgent(agent); + if (cacheConfig != null) { + app.contextCacheConfig(cacheConfig); + } + return Runner.builder() + .app(app.build()) + .sessionService(new InMemorySessionService()) + .artifactService(new InMemoryArtifactService()) + .build(); + } + + private static String newSession(Runner runner) { + return runner.sessionService().createSession(runner.appName(), USER_ID).blockingGet().id(); + } + + /** Returns the name of the thread each event of one turn was emitted on. */ + private static ImmutableList eventThreads( + Runner runner, String sessionId, String text, RunConfig runConfig) { + return ImmutableList.copyOf( + runner + .runAsync(USER_ID, sessionId, Content.fromParts(Part.fromText(text)), runConfig) + .map(event -> Thread.currentThread().getName()) + .toList() + .blockingGet()); + } + + @CanIgnoreReturnValue + private static ImmutableList run(Runner runner, String sessionId, String text) { + return run(runner, sessionId, text, RunConfig.builder().build()); + } + + @CanIgnoreReturnValue + private static ImmutableList run( + Runner runner, String sessionId, String text, RunConfig runConfig) { + return ImmutableList.copyOf( + runner + .runAsync(USER_ID, sessionId, Content.fromParts(Part.fromText(text)), runConfig) + .toList() + .blockingGet()); + } + + private static LlmRequest.Builder request(Content... contents) { + return LlmRequest.builder() + .model(MODEL) + .contents(ImmutableList.copyOf(contents)) + .config(CONFIG) + .cacheConfig(CACHE_CONFIG); + } + + private static Instant inOneHour() { + return NOW.plus(Duration.ofHours(1)); + } + + private static Content text(String role, String text) { + return Content.builder().role(role).parts(Part.fromText(text)).build(); + } + + private static Tool functionTool(String... names) { + return Tool.builder() + .functionDeclarations( + Arrays.stream(names) + .map(name -> FunctionDeclaration.builder().name(name).description(name).build()) + .collect(toImmutableList())) + .build(); + } + + private static ImmutableList texts(Content... contents) { + return Arrays.stream(contents) + .map(content -> content.parts().get().get(0).text().get()) + .collect(toImmutableList()); + } + + private static ImmutableList texts(JsonNode contents) { + return stream(contents) + .flatMap(content -> stream(content.get("parts"))) + .map(part -> part.path("text").asText()) + .collect(toImmutableList()); + } + + private enum Kind { + GENERATE, + CREATE_CACHE, + DELETE_CACHE + } + + /** + * Answers the Gemini API over HTTP. Each model call gets the next numbered answer, and each cache + * create returns a new numbered cache that expires in an hour, unless told to fail or to omit the + * expiry. + */ + private static final class FakeGeminiApi implements Interceptor { + private final List calls = new ArrayList<>(); + private final List expireTimes = new ArrayList<>(); + private final Deque createFailures = new ArrayDeque<>(); + private final Deque deleteFailures = new ArrayDeque<>(); + private @Nullable String nextCreateBody = null; + private boolean omitExpireTime = false; + private Duration cacheCallDelay = Duration.ZERO; + private int answers = 0; + + Client client() { + OkHttpClient httpClient = new OkHttpClient.Builder().addInterceptor(this).build(); + return Client.builder() + .apiKey("test-api-key") + .vertexAI(false) + .clientOptions(ClientOptions.builder().customHttpClient(httpClient).build()) + .build(); + } + + synchronized void failNextCreate(int code) { + createFailures.add(code); + } + + synchronized void failNextDelete(int code) { + deleteFailures.add(code); + } + + synchronized void answerNextCreateWith(String body) { + nextCreateBody = body; + } + + synchronized void omitExpireTime() { + omitExpireTime = true; + } + + synchronized void delayCacheCalls(Duration delay) { + cacheCallDelay = delay; + } + + /** Returns the expiry reported for the {@code n}-th created cache, counting from 1. */ + synchronized Instant expireTime(int n) { + return expireTimes.get(n - 1); + } + + synchronized ImmutableList kinds() { + return calls.stream().map(Call::kind).collect(toImmutableList()); + } + + synchronized ImmutableList paths(Kind kind) { + return calls.stream() + .filter(call -> call.kind() == kind) + .map(Call::path) + .collect(toImmutableList()); + } + + synchronized ImmutableList bodies(Kind kind) { + return calls.stream() + .filter(call -> call.kind() == kind) + .map(Call::body) + .collect(toImmutableList()); + } + + synchronized ImmutableList headers(Kind kind) { + return calls.stream() + .filter(call -> call.kind() == kind) + .map(Call::headers) + .collect(toImmutableList()); + } + + @Override + public synchronized Response intercept(Chain chain) throws IOException { + Request request = chain.request(); + String path = request.url().encodedPath(); + JsonNode body = MissingNode.getInstance(); + if (request.body() != null) { + Buffer buffer = new Buffer(); + request.body().writeTo(buffer); + body = objectMapper.readTree(buffer.readUtf8()); + } + Kind kind; + int code = 200; + String responseBody; + if (path.endsWith(":generateContent")) { + kind = Kind.GENERATE; + answers++; + responseBody = answer("Answer " + answers).toJson(); + } else if (path.endsWith(":streamGenerateContent")) { + kind = Kind.GENERATE; + answers++; + responseBody = + "data: " + + partialAnswer("Answer ").toJson() + + "\n\ndata: " + + answer(String.valueOf(answers)).toJson() + + "\n\n"; + } else if (request.method().equals("POST") && path.endsWith("/cachedContents")) { + kind = Kind.CREATE_CACHE; + sleep(cacheCallDelay); + if (!createFailures.isEmpty()) { + code = createFailures.remove(); + responseBody = error(code); + } else if (nextCreateBody != null) { + responseBody = nextCreateBody; + nextCreateBody = null; + } else { + Instant expireTime = Instant.now().plus(Duration.ofHours(1)); + expireTimes.add(expireTime); + String name = "cachedContents/cache-" + expireTimes.size(); + responseBody = + omitExpireTime + ? String.format("{\"name\": \"%s\"}", name) + : String.format("{\"name\": \"%s\", \"expireTime\": \"%s\"}", name, expireTime); + } + } else if (request.method().equals("DELETE") && path.contains("/cachedContents/")) { + kind = Kind.DELETE_CACHE; + sleep(cacheCallDelay); + code = deleteFailures.isEmpty() ? 200 : deleteFailures.remove(); + responseBody = code == 200 ? "{}" : error(code); + } else { + throw new IOException("Unexpected request: " + request.method() + " " + path); + } + calls.add(new Call(kind, path, body, request.headers())); + return new Response.Builder() + .request(request) + .protocol(Protocol.HTTP_1_1) + .code(code) + .message("fake") + .body(ResponseBody.create(responseBody, MediaType.get("application/json"))) + .build(); + } + + private static GenerateContentResponse answer(String text) { + return GenerateContentResponse.builder() + .candidates( + Candidate.builder() + .content(text("model", text)) + .finishReason(new FinishReason(FinishReason.Known.STOP)) + .build()) + .usageMetadata( + GenerateContentResponseUsageMetadata.builder() + .promptTokenCount(5000) + .candidatesTokenCount(10) + .totalTokenCount(5010) + .build()) + .build(); + } + + private static GenerateContentResponse partialAnswer(String text) { + return GenerateContentResponse.builder() + .candidates(Candidate.builder().content(text("model", text)).build()) + .build(); + } + + private static void sleep(Duration duration) throws IOException { + try { + Thread.sleep(duration.toMillis()); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IOException(e); + } + } + + private static String error(int code) { + return "{\"error\": {\"code\": " + code + ", \"message\": \"fake\"}}"; + } + + private record Call(Kind kind, String path, JsonNode body, Headers headers) {} + } +} diff --git a/core/src/test/java/com/google/adk/models/LlmRequestTest.java b/core/src/test/java/com/google/adk/models/LlmRequestTest.java index 56c2debd7..953406845 100644 --- a/core/src/test/java/com/google/adk/models/LlmRequestTest.java +++ b/core/src/test/java/com/google/adk/models/LlmRequestTest.java @@ -19,6 +19,7 @@ import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertThrows; +import com.google.adk.agents.ContextCacheConfig; import com.google.adk.tools.BaseTool; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; @@ -27,6 +28,7 @@ import com.google.genai.types.LiveConnectConfig; import com.google.genai.types.Part; import com.google.genai.types.Schema; +import java.time.Duration; import java.util.Optional; import org.junit.Test; import org.junit.runner.RunWith; @@ -308,4 +310,35 @@ public void getSystemInstructions_whenPresent_returnsList() { .containsExactly(instruction1 + "\n\n" + instruction2) .inOrder(); } + + @Test + public void cacheFields_surviveToBuilderAndStayOutOfJson() { + ContextCacheConfig cacheConfig = new ContextCacheConfig(5, Duration.ofMinutes(10), 1024); + CacheMetadata cacheMetadata = + CacheMetadata.builder().fingerprint("abc123").contentsCount(2).build(); + + LlmRequest request = + LlmRequest.builder() + .model("gemini-2.5-flash") + .cacheConfig(cacheConfig) + .cacheMetadata(cacheMetadata) + .cacheableContentsTokenCount(4096) + .build() + .toBuilder() + .build(); + + assertThat(request.cacheConfig()).hasValue(cacheConfig); + assertThat(request.cacheMetadata()).hasValue(cacheMetadata); + assertThat(request.cacheableContentsTokenCount()).hasValue(4096); + assertThat(request.toJson()).doesNotContain("cache"); + } + + @Test + public void cacheFields_emptyByDefault() { + LlmRequest request = LlmRequest.builder().build(); + + assertThat(request.cacheConfig()).isEmpty(); + assertThat(request.cacheMetadata()).isEmpty(); + assertThat(request.cacheableContentsTokenCount()).isEmpty(); + } } diff --git a/core/src/test/java/com/google/adk/models/LlmResponseTest.java b/core/src/test/java/com/google/adk/models/LlmResponseTest.java index 44b16854b..c930f73b5 100644 --- a/core/src/test/java/com/google/adk/models/LlmResponseTest.java +++ b/core/src/test/java/com/google/adk/models/LlmResponseTest.java @@ -144,6 +144,26 @@ public void testSerializationAndDeserialization_optionalFieldsEmpty() assertThat(deserializedResponse.usageMetadata()).isEmpty(); } + @Test + public void testSerializationAndDeserialization_withCacheMetadata() + throws JsonProcessingException { + CacheMetadata cacheMetadata = + CacheMetadata.builder().fingerprint("abc123").contentsCount(2).build(); + LlmResponse originalResponse = + LlmResponse.builder() + .content(createSampleContent("hello")) + .cacheMetadata(cacheMetadata) + .build(); + + String json = originalResponse.toJson(); + + assertThat(objectMapper.readTree(json).get("cacheMetadata").get("fingerprint").asText()) + .isEqualTo("abc123"); + LlmResponse deserializedResponse = LlmResponse.fromJsonString(json, LlmResponse.class); + assertThat(deserializedResponse).isEqualTo(originalResponse); + assertThat(deserializedResponse.cacheMetadata()).hasValue(cacheMetadata); + } + @Test public void testSerializationAndDeserialization_withTranscriptions() throws JsonProcessingException { diff --git a/examples/README.md b/examples/README.md index 7ab11f6b1..65768a1b9 100644 --- a/examples/README.md +++ b/examples/README.md @@ -13,6 +13,8 @@ contains more complete end-to-end samples for customers to use or modify. - [`a2a_server`](a2a_server/README.md): an ADK agent served over A2A with Quarkus. - [`configagent`](configagent/README.md): agents defined in YAML config files. +- [`contextcaching`](contextcaching/README.md): Gemini context caching that + reuses and refreshes a cache across turns. - [`helloworld`](helloworld/README.md): a minimal agent with dice-rolling and prime-checking tools. - [`mcpfilesystem`](mcpfilesystem/README.md): an agent that uses the MCP diff --git a/examples/contextcaching/ContextCachingAgent.java b/examples/contextcaching/ContextCachingAgent.java new file mode 100644 index 000000000..2ee3bb495 --- /dev/null +++ b/examples/contextcaching/ContextCachingAgent.java @@ -0,0 +1,126 @@ +/* + * 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.example.contextcaching; + +import com.google.adk.agents.ContextCacheConfig; +import com.google.adk.agents.LlmAgent; +import com.google.adk.apps.App; +import com.google.common.collect.ImmutableList; +import java.time.Duration; + +/** A weather assistant whose large, shared instruction makes it a good fit for context caching. */ +public final class ContextCachingAgent { + + static final String DEFAULT_MODEL = "gemini-3.8-flash"; + + /** + * The shortest TTL Vertex AI accepts, so the demo can let the cache expire and create it again. + */ + static final Duration CACHE_TTL = Duration.ofMinutes(1); + + /** + * The system instruction, which is the bulk of the cached prefix. Explicit caching only starts + * once that prefix reaches the model's minimum (2048 tokens for Gemini 2.5, 4096 for Gemini 3), + * so generated station rows pad it past that. + */ + private static final String WEATHER_REFERENCE = buildWeatherReference(); + + private ContextCachingAgent() {} + + /** Returns an app that enables context caching for all its agents. */ + public static App createApp(String model) { + LlmAgent agent = + LlmAgent.builder() + .name("weather_assistant") + .description("Answers weather questions from a large shared reference.") + .model(model) + .instruction(WEATHER_REFERENCE) + .build(); + return App.builder() + .name("context_caching_demo") + .rootAgent(agent) + .contextCacheConfig( + new ContextCacheConfig(/* cacheIntervals= */ 10, CACHE_TTL, /* minTokens= */ 0)) + .build(); + } + + private static String buildWeatherReference() { + StringBuilder reference = new StringBuilder(); + reference.append( + "You are a weather assistant. Answer questions using ONLY the weather reference below," + + " keep answers to one or two sentences, and cite the section or city you used.\n\n"); + + ImmutableList guidelines = + ImmutableList.of( + "Forecasting basics: a forecast combines current observations, numerical model" + + " guidance, and local climatology; always state the valid time window and your" + + " confidence.", + "Wind chill: how cold the air feels once wind is accounted for. It is only defined for" + + " temperatures at or below 10C and wind above 5 km/h.", + "Heat index: the apparent temperature, combining air temperature and humidity. It is" + + " meaningful above about 27C when relative humidity is high.", + "Precipitation types: rain, drizzle, sleet, freezing rain, snow, and hail are" + + " distinguished by the temperature profile between the cloud base and the" + + " ground.", + "Storm safety: during a thunderstorm, move indoors, avoid open fields and tall isolated" + + " trees, and stay off corded electronics until 30 minutes after the last" + + " thunder.", + "Units: temperatures are in degrees Celsius, wind speed in km/h, and precipitation in" + + " millimeters unless a value states otherwise."); + for (int i = 0; i < guidelines.size(); i++) { + reference.append("Section ").append(i + 1).append(": ").append(guidelines.get(i)); + reference.append('\n'); + } + + reference.append("\nClimate normals for named cities:\n"); + ImmutableList cities = + ImmutableList.of( + "Marisol: coastal and mild; January average 12C, July average 24C, annual rainfall" + + " 640 mm.", + "Fjordheim: cold maritime; January average -6C, July average 15C, annual rainfall" + + " 900 mm.", + "Solara: hot desert; January average 18C, July average 41C, annual rainfall 90 mm.", + "Verdant: temperate rainforest; January average 7C, July average 19C, annual rainfall" + + " 2400 mm.", + "Highpoint: alpine; January average -9C, July average 12C, annual rainfall 1100 mm.", + "Puerto Brisa: tropical; January average 26C, July average 29C, annual rainfall" + + " 1800 mm.", + "Windgate: windy plains; January average -2C, July average 27C, annual rainfall" + + " 520 mm.", + "Frosthollow: subarctic; January average -22C, July average 16C, annual rainfall" + + " 380 mm."); + cities.forEach(city -> reference.append(" ").append(city).append('\n')); + + reference.append("\nAutomated station normals:\n"); + ImmutableList windDirections = + ImmutableList.of("N", "NE", "E", "SE", "S", "SW", "W", "NW"); + for (int i = 1; i <= 100; i++) { + reference.append( + String.format( + " Station S%d (sector %c%d): January average %dC, July average %dC, annual rainfall" + + " %d mm, prevailing wind %s at %d km/h.\n", + i, + (char) ('A' + i % 8), + i % 12, + -12 + (i * 7) % 34, + 14 + (i * 3) % 22, + 250 + (i * 37) % 2300, + windDirections.get(i % 8), + 6 + (i * 5) % 38)); + } + return reference.toString(); + } +} diff --git a/examples/contextcaching/ContextCachingRun.java b/examples/contextcaching/ContextCachingRun.java new file mode 100644 index 000000000..494e21f89 --- /dev/null +++ b/examples/contextcaching/ContextCachingRun.java @@ -0,0 +1,111 @@ +/* + * 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.example.contextcaching; + +import com.google.adk.apps.App; +import com.google.adk.events.Event; +import com.google.adk.models.CacheMetadata; +import com.google.adk.runner.Runner; +import com.google.common.collect.ImmutableList; +import com.google.genai.types.Content; +import com.google.genai.types.Part; +import java.time.Duration; +import org.jspecify.annotations.Nullable; + +/** + * Runs several turns of one session against {@link ContextCachingAgent} and prints what the cache + * does on each turn. It pauses once so the cache expires, and the next turn creates a new one. + * Exact turn-by-turn output depends on model latency, since a slow turn can outlast the TTL. + */ +public final class ContextCachingRun { + + private static final ImmutableList PROMPTS = + ImmutableList.of( + "In one sentence, what should I do during a thunderstorm?", + "How is wind chill defined, and when does it apply?", + "What units does the reference use for wind speed?", + "What is the average annual rainfall listed for Verdant?", + "Which named city is the driest?", + "What is the July average temperature for Puerto Brisa?"); + + /** Zero-based turn after which the demo waits for the cache to expire. */ + private static final int EXPIRY_AFTER_TURN = 2; + + private static final String USER_ID = "demo-user"; + + private ContextCachingRun() {} + + public static void main(String[] args) throws InterruptedException { + String model = args.length > 0 ? args[0] : ContextCachingAgent.DEFAULT_MODEL; + App app = ContextCachingAgent.createApp(model); + Runner runner = Runner.builder().app(app).build(); + String sessionId = + runner.sessionService().createSession(app.name(), USER_ID).blockingGet().id(); + + System.out.printf( + "Context caching demo on %s, cache TTL %ds.%n", + model, ContextCachingAgent.CACHE_TTL.toSeconds()); + + String activeCacheName = null; + for (int turn = 0; turn < PROMPTS.size(); turn++) { + String prompt = PROMPTS.get(turn); + System.out.printf("%n===== Turn %d: %s =====%n", turn + 1, prompt); + Event answer = + runner + .runAsync(USER_ID, sessionId, Content.fromParts(Part.fromText(prompt))) + .blockingLast(); + System.out.println("assistant> " + answer.stringifyContent().strip()); + + CacheMetadata cache = answer.cacheMetadata().orElse(null); + if (cache != null) { + System.out.printf( + " cache: %s | cachedContents=%d invocationsUsed=%d%n", + describe(cache, activeCacheName), + cache.contentsCount(), + cache.invocationsUsed().orElse(0)); + activeCacheName = cache.cacheName().orElse(activeCacheName); + } + answer + .usageMetadata() + .ifPresent( + usage -> + System.out.printf( + " prompt tokens: %d, served from cache: %d%n", + usage.promptTokenCount().orElse(0), + usage.cachedContentTokenCount().orElse(0))); + + if (turn == EXPIRY_AFTER_TURN) { + Duration wait = ContextCachingAgent.CACHE_TTL.plusSeconds(2); + System.out.printf("%n... waiting %ds so the cache expires ...%n", wait.toSeconds()); + Thread.sleep(wait.toMillis()); + } + } + } + + /** Describes a turn's cache relative to the cache the previous turns used. */ + private static String describe(CacheMetadata cache, @Nullable String previousCacheName) { + String name = cache.cacheName().orElse(null); + if (name == null) { + return previousCacheName == null + ? "no cache yet (fingerprint only)" + : "no active cache (fingerprint only)"; + } + if (previousCacheName == null) { + return "CREATED " + name; + } + return name.equals(previousCacheName) ? "REUSED " + name : "RE-CREATED " + name; + } +} diff --git a/examples/contextcaching/README.md b/examples/contextcaching/README.md new file mode 100644 index 000000000..147534be7 --- /dev/null +++ b/examples/contextcaching/README.md @@ -0,0 +1,54 @@ +# Context Caching Sample + +This sample shows explicit context caching with the Google ADK. A weather +assistant answers several questions in one session. Its large system instruction +is shared by every turn, so the app caches it on the Gemini backend with +`App.contextCacheConfig`, and later turns reference the cached instruction +instead of resending it. + +The runner prints what the cache does on each turn: + +- Turn 1 has no cache yet. It only records a fingerprint of the cacheable + prefix. +- Once the prefix is known to be stable, a cache is CREATED. Later turns REUSE + it, and `invocationsUsed` grows. +- After the third turn the runner waits for the cache's one-minute TTL, the + shortest Vertex AI accepts, to run out. The next turn RE-CREATES the cache, + usually with more cached contents. +- Each turn also prints its prompt token count and how many of those tokens + the cache served. + +Exact turn-by-turn output depends on model latency. A slow turn can outlast the +one-minute TTL, so a later turn may print an extra RE-CREATED. + +Caching only starts once the cached prefix reaches the model's minimum size: +2048 tokens for Gemini 2.5 and 4096 tokens for Gemini 3. + +## Project Layout + +``` +├── ContextCachingAgent.java // Agent, app and caching configuration +├── ContextCachingRun.java // Console runner entry point +├── pom.xml // Maven configuration and exec main class +└── README.md // This file +``` + +## Prerequisites + +- Java 17+ +- Maven 3.9+ +- Credentials for the Gemini API (`GOOGLE_API_KEY`) or Vertex AI + (`GOOGLE_GENAI_USE_VERTEXAI=true`, `GOOGLE_CLOUD_PROJECT` and + `GOOGLE_CLOUD_LOCATION`) + +## Build and Run + +```bash +mvn clean compile exec:java +``` + +The sample uses `gemini-3.8-flash`. To try another model, pass its name: + +```bash +mvn compile exec:java -Dexec.args="" +``` diff --git a/examples/contextcaching/pom.xml b/examples/contextcaching/pom.xml new file mode 100644 index 000000000..dedbb3197 --- /dev/null +++ b/examples/contextcaching/pom.xml @@ -0,0 +1,104 @@ + + + + 4.0.0 + + + com.google.adk + google-adk-examples + 1.11.1-SNAPSHOT + .. + + + com.google.adk.samples + google-adk-sample-contextcaching + Google ADK - Sample - Context Caching + + A sample application demonstrating context caching with the Google ADK, + runnable via com.example.contextcaching.ContextCachingRun. + + jar + + + UTF-8 + 17 + + com.example.contextcaching.ContextCachingRun + ${project.version} + + + + + com.google.adk + google-adk + ${google-adk.version} + + + + + + + org.apache.maven.plugins + maven-compiler-plugin + 3.13.0 + + ${java.version} + ${java.version} + true + + + + org.codehaus.mojo + build-helper-maven-plugin + 3.6.0 + + + add-source + generate-sources + + add-source + + + + . + + + + + + + org.apache.maven.plugins + maven-source-plugin + + + **/*.jar + target/** + + + + + org.codehaus.mojo + exec-maven-plugin + 3.2.0 + + ${exec.mainClass} + runtime + + + + + diff --git a/examples/pom.xml b/examples/pom.xml index cfb1517a0..eb3fd8a9d 100644 --- a/examples/pom.xml +++ b/examples/pom.xml @@ -34,6 +34,7 @@ a2a_basic a2a_server configagent + contextcaching helloworld mcpfilesystem