diff --git a/core/src/main/java/com/google/adk/agents/InvocationContext.java b/core/src/main/java/com/google/adk/agents/InvocationContext.java index f3a3cf3b0..047fb0b1c 100644 --- a/core/src/main/java/com/google/adk/agents/InvocationContext.java +++ b/core/src/main/java/com/google/adk/agents/InvocationContext.java @@ -21,6 +21,8 @@ import com.google.adk.artifacts.BaseArtifactService; import com.google.adk.memory.BaseMemoryService; import com.google.adk.models.LlmCallsLimitExceededException; +import com.google.adk.platform.TimeProvider; +import com.google.adk.platform.UuidProvider; import com.google.adk.plugins.Plugin; import com.google.adk.plugins.PluginManager; import com.google.adk.sessions.BaseSessionService; @@ -28,10 +30,10 @@ import com.google.adk.summarizer.EventsCompactionConfig; import com.google.errorprone.annotations.CanIgnoreReturnValue; import com.google.genai.types.Content; +import java.time.Instant; import java.util.Map; import java.util.Objects; import java.util.Optional; -import java.util.UUID; import java.util.concurrent.ConcurrentHashMap; import org.jspecify.annotations.Nullable; @@ -52,6 +54,8 @@ public class InvocationContext { @Nullable private final ContextCacheConfig contextCacheConfig; private final InvocationCostManager invocationCostManager; private final Map callbackContextData; + private final TimeProvider timeProvider; + private final UuidProvider uuidProvider; @Nullable private String branch; private BaseAgent agent; @@ -78,6 +82,8 @@ protected InvocationContext(Builder builder) { // invocation invocation so that Plugins can access the same data it during the invocation // across all types of callbacks. this.callbackContextData = builder.callbackContextData; + this.timeProvider = builder.timeProvider; + this.uuidProvider = builder.uuidProvider; } /** Returns a new {@link Builder} for creating {@link InvocationContext} instances. */ @@ -192,9 +198,34 @@ public String userId() { return session.userId(); } + /** Returns the {@link TimeProvider} for this invocation. */ + public TimeProvider timeProvider() { + return timeProvider; + } + + /** Returns the {@link UuidProvider} for this invocation. */ + public UuidProvider uuidProvider() { + return uuidProvider; + } + + /** Returns the current time from this invocation's {@link TimeProvider}. */ + public Instant now() { + return timeProvider.now(); + } + + /** Returns a new unique identifier from this invocation's {@link UuidProvider}. */ + public String newUuid() { + return uuidProvider.newUuid(); + } + /** Generates a new unique ID for an invocation context. */ public static String newInvocationContextId() { - return "e-" + UUID.randomUUID(); + return newInvocationContextId(UuidProvider.SYSTEM); + } + + /** Generates a new unique ID for an invocation context using the given {@link UuidProvider}. */ + public static String newInvocationContextId(UuidProvider uuidProvider) { + return "e-" + uuidProvider.newUuid(); } /** @@ -275,6 +306,8 @@ private Builder(InvocationContext context) { // invocation invocation so that Plugins can access the same data it during the invocation // across all types of callbacks. this.callbackContextData = context.callbackContextData; + this.timeProvider = context.timeProvider; + this.uuidProvider = context.uuidProvider; } private BaseSessionService sessionService; @@ -294,6 +327,8 @@ private Builder(InvocationContext context) { @Nullable private ContextCacheConfig contextCacheConfig; private InvocationCostManager invocationCostManager = new InvocationCostManager(); private Map callbackContextData = new ConcurrentHashMap<>(); + private TimeProvider timeProvider = TimeProvider.SYSTEM; + private UuidProvider uuidProvider = UuidProvider.SYSTEM; /** * Sets the session service for managing session state. @@ -475,6 +510,30 @@ public Builder callbackContextData(Map callbackContextData) { return this; } + /** + * Sets the time provider for the invocation. Defaults to {@link TimeProvider#SYSTEM}. + * + * @param timeProvider the provider for the current time. + * @return this builder instance for chaining. + */ + @CanIgnoreReturnValue + public Builder timeProvider(TimeProvider timeProvider) { + this.timeProvider = timeProvider; + return this; + } + + /** + * Sets the UUID provider for the invocation. Defaults to {@link UuidProvider#SYSTEM}. + * + * @param uuidProvider the provider for new unique identifiers. + * @return this builder instance for chaining. + */ + @CanIgnoreReturnValue + public Builder uuidProvider(UuidProvider uuidProvider) { + this.uuidProvider = uuidProvider; + return this; + } + /** * Builds the {@link InvocationContext} instance. * @@ -531,7 +590,9 @@ public boolean equals(Object o) { && Objects.equals(eventsCompactionConfig, that.eventsCompactionConfig) && Objects.equals(contextCacheConfig, that.contextCacheConfig) && Objects.equals(invocationCostManager, that.invocationCostManager) - && Objects.equals(callbackContextData, that.callbackContextData); + && Objects.equals(callbackContextData, that.callbackContextData) + && Objects.equals(timeProvider, that.timeProvider) + && Objects.equals(uuidProvider, that.uuidProvider); } @Override @@ -553,6 +614,8 @@ public int hashCode() { eventsCompactionConfig, contextCacheConfig, invocationCostManager, - callbackContextData); + callbackContextData, + timeProvider, + uuidProvider); } } 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 3b28761a1..9796d82c7 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 @@ -438,7 +438,7 @@ private Flowable runOneStep(Context spanContext, InvocationContext contex final Event mutableEventTemplate = Event.builder() - .id(Event.generateEventId()) + .id(context.newUuid()) .invocationId(context.invocationId()) .author(context.agent().name()) .branch(context.branch().orElse(null)) @@ -453,7 +453,7 @@ private Flowable runOneStep(Context spanContext, InvocationContext contex .doFinally( () -> { String oldId = mutableEventTemplate.id(); - String newId = Event.generateEventId(); + String newId = context.newUuid(); logger.debug("Resetting event ID from {} to {}", oldId, newId); mutableEventTemplate.setId(newId); }) @@ -461,7 +461,7 @@ private Flowable runOneStep(Context spanContext, InvocationContext contex event -> { // Update event ID for the new resulting events String oldId = event.id(); - String newId = Event.generateEventId(); + String newId = context.newUuid(); logger.debug("Resetting event ID from {} to {}", oldId, newId); event = event.toBuilder().id(newId).build(); Flowable postProcessedEvents = Flowable.just(event); @@ -555,7 +555,7 @@ public Flowable runLive(InvocationContext invocationContext) { return Flowable.empty(); } - String eventIdForSendData = Event.generateEventId(); + String eventIdForSendData = invocationContext.newUuid(); LlmAgent agent = (LlmAgent) invocationContext.agent(); BaseLlm llm = agent.resolvedModel().model().isPresent() @@ -647,7 +647,7 @@ public void onError(Throwable e) { .flatMap( llmResponse -> { Event baseEventForThisLlmResponse = - liveEventBuilderTemplate.id(Event.generateEventId()).build(); + liveEventBuilderTemplate.id(invocationContext.newUuid()).build(); return postprocess( invocationContext, baseEventForThisLlmResponse, @@ -727,7 +727,7 @@ private Flowable buildPostprocessingEvents( } Event modelResponseEvent = - buildModelResponseEvent(baseEventForLlmResponse, llmRequest, updatedResponse); + buildModelResponseEvent(context, baseEventForLlmResponse, llmRequest, updatedResponse); if (modelResponseEvent.functionCalls().isEmpty()) { return processorEvents.concatWith(Flowable.just(modelResponseEvent)); } @@ -772,9 +772,13 @@ private void traceCallLlm( } private Event buildModelResponseEvent( - Event baseEventForLlmResponse, LlmRequest llmRequest, LlmResponse llmResponse) { + InvocationContext context, + Event baseEventForLlmResponse, + LlmRequest llmRequest, + LlmResponse llmResponse) { Event.Builder eventBuilder = baseEventForLlmResponse.toBuilder() + .timestamp(context.now().toEpochMilli()) .content(llmResponse.content().orElse(null)) .partial(llmResponse.partial().orElse(null)) .errorCode(llmResponse.errorCode().orElse(null)) @@ -794,7 +798,7 @@ private Event buildModelResponseEvent( logger.debug("event: {} functionCalls: {}", event, event.functionCalls()); if (!event.functionCalls().isEmpty()) { - Functions.populateClientFunctionCallId(event); + Functions.populateClientFunctionCallId(event, context.uuidProvider()); Set longRunningToolIds = Functions.getLongRunningFunctionCalls(event.functionCalls(), llmRequest.tools()); logger.debug("longRunningToolIds: {}", longRunningToolIds); diff --git a/core/src/main/java/com/google/adk/flows/llmflows/Functions.java b/core/src/main/java/com/google/adk/flows/llmflows/Functions.java index 8c60ebf76..2b2dbff49 100644 --- a/core/src/main/java/com/google/adk/flows/llmflows/Functions.java +++ b/core/src/main/java/com/google/adk/flows/llmflows/Functions.java @@ -29,6 +29,7 @@ import com.google.adk.events.Event; import com.google.adk.events.EventActions; import com.google.adk.events.ToolConfirmation; +import com.google.adk.platform.UuidProvider; import com.google.adk.telemetry.Instrumentation; import com.google.adk.telemetry.Instrumentation.ToolExecution; import com.google.adk.telemetry.Tracing; @@ -58,7 +59,6 @@ import java.util.Map; import java.util.Optional; import java.util.Set; -import java.util.UUID; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -75,7 +75,12 @@ public final class Functions { /** Generates a unique ID for a function call. */ public static String generateClientFunctionCallId() { - return AF_FUNCTION_CALL_ID_PREFIX + UUID.randomUUID(); + return generateClientFunctionCallId(UuidProvider.SYSTEM); + } + + /** Generates a unique ID for a function call using the given {@link UuidProvider}. */ + public static String generateClientFunctionCallId(UuidProvider uuidProvider) { + return AF_FUNCTION_CALL_ID_PREFIX + uuidProvider.newUuid(); } /** @@ -87,6 +92,18 @@ public static String generateClientFunctionCallId() { * @param modelResponseEvent The event potentially containing function calls. */ public static void populateClientFunctionCallId(Event modelResponseEvent) { + populateClientFunctionCallId(modelResponseEvent, UuidProvider.SYSTEM); + } + + /** + * Populates missing function call IDs in the provided event's content using the given {@link + * UuidProvider}. + * + * @param modelResponseEvent The event potentially containing function calls. + * @param uuidProvider The provider used to mint new function call IDs. + */ + public static void populateClientFunctionCallId( + Event modelResponseEvent, UuidProvider uuidProvider) { Optional originalContentOptional = modelResponseEvent.content(); if (originalContentOptional.isEmpty()) { return; @@ -104,7 +121,7 @@ public static void populateClientFunctionCallId(Event modelResponseEvent) { FunctionCall functionCall = part.functionCall().get(); if (functionCall.id().isEmpty() || functionCall.id().get().isEmpty()) { FunctionCall updatedFunctionCall = - functionCall.toBuilder().id(generateClientFunctionCallId()).build(); + functionCall.toBuilder().id(generateClientFunctionCallId(uuidProvider)).build(); newParts.add(part.toBuilder().functionCall(updatedFunctionCall).build()); modified = true; } else { @@ -169,7 +186,7 @@ public static Maybe handleFunctionCalls( return Maybe.empty(); } Optional maybeMergedEvent = - Functions.mergeParallelFunctionResponseEvents(events); + Functions.mergeParallelFunctionResponseEvents(invocationContext, events); if (maybeMergedEvent.isEmpty()) { return Maybe.empty(); } @@ -233,7 +250,8 @@ public static Maybe handleFunctionCallsLive( if (events.isEmpty()) { return Maybe.empty(); } - return Maybe.fromOptional(Functions.mergeParallelFunctionResponseEvents(events)); + return Maybe.fromOptional( + Functions.mergeParallelFunctionResponseEvents(invocationContext, events)); }); } @@ -492,7 +510,7 @@ private static Maybe processFunctionResult( } private static Optional mergeParallelFunctionResponseEvents( - List functionResponseEvents) { + InvocationContext invocationContext, List functionResponseEvents) { if (functionResponseEvents.isEmpty()) { return Optional.empty(); } @@ -516,7 +534,7 @@ private static Optional mergeParallelFunctionResponseEvents( return Optional.of( Event.builder() - .id(Event.generateEventId()) + .id(invocationContext.newUuid()) .invocationId(baseEvent.invocationId()) .author(baseEvent.author()) .branch(baseEvent.branch().orElse(null)) @@ -667,7 +685,8 @@ private static Event buildResponseEvent( .build(); return Event.builder() - .id(Event.generateEventId()) + .id(invocationContext.newUuid()) + .timestamp(invocationContext.now().toEpochMilli()) .invocationId(invocationContext.invocationId()) .author(invocationContext.agent().name()) .branch(invocationContext.branch().orElse(null)) @@ -712,7 +731,7 @@ public static Optional generateRequestConfirmationEvent( functionCallsById.get(entry.getKey()), "toolConfirmation", entry.getValue())) - .id(generateClientFunctionCallId()) + .id(generateClientFunctionCallId(invocationContext.uuidProvider())) .build(); longRunningToolIds.add(requestConfirmationFunctionCall.id().get()); @@ -728,6 +747,8 @@ public static Optional generateRequestConfirmationEvent( return Optional.of( Event.builder() + .id(invocationContext.newUuid()) + .timestamp(invocationContext.now().toEpochMilli()) .invocationId(invocationContext.invocationId()) .author(invocationContext.agent().name()) .branch(invocationContext.branch().orElse(null)) diff --git a/core/src/main/java/com/google/adk/flows/llmflows/OutputSchema.java b/core/src/main/java/com/google/adk/flows/llmflows/OutputSchema.java index d1f322f18..f1543c81e 100644 --- a/core/src/main/java/com/google/adk/flows/llmflows/OutputSchema.java +++ b/core/src/main/java/com/google/adk/flows/llmflows/OutputSchema.java @@ -109,7 +109,8 @@ public static Optional getStructuredModelResponse(Event functionResponse public static Event createFinalModelResponseEvent( InvocationContext context, String jsonResponse) { return Event.builder() - .id(Event.generateEventId()) + .id(context.newUuid()) + .timestamp(context.now().toEpochMilli()) .invocationId(context.invocationId()) .author(context.agent().name()) .branch(context.branch().orElse(null)) diff --git a/core/src/main/java/com/google/adk/platform/TimeProvider.java b/core/src/main/java/com/google/adk/platform/TimeProvider.java new file mode 100644 index 000000000..9445c5146 --- /dev/null +++ b/core/src/main/java/com/google/adk/platform/TimeProvider.java @@ -0,0 +1,35 @@ +/* + * 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.platform; + +import java.time.Instant; + +/** + * Supplies the current time for ADK-generated timestamps. + * + *

The default {@link #SYSTEM} provider reads the wall clock. Integrations that need custom + * timestamps can install a custom provider on an invocation; see {@code InvocationContext}. + */ +@FunctionalInterface +public interface TimeProvider { + + /** A provider backed by the system wall clock ({@link Instant#now()}). */ + TimeProvider SYSTEM = Instant::now; + + /** Returns the current time. */ + Instant now(); +} diff --git a/core/src/main/java/com/google/adk/platform/UuidProvider.java b/core/src/main/java/com/google/adk/platform/UuidProvider.java new file mode 100644 index 000000000..ef38be0f0 --- /dev/null +++ b/core/src/main/java/com/google/adk/platform/UuidProvider.java @@ -0,0 +1,35 @@ +/* + * 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.platform; + +import java.util.UUID; + +/** + * Supplies new unique identifiers for ADK-generated IDs (events, invocations, function calls). + * + *

The default {@link #SYSTEM} provider returns random UUIDs. Integrations that need customized + * identifiers can install * a custom provider on an invocation; see {@code InvocationContext}. + */ +@FunctionalInterface +public interface UuidProvider { + + /** A provider backed by {@link UUID#randomUUID()}. */ + UuidProvider SYSTEM = () -> UUID.randomUUID().toString(); + + /** Returns a new unique identifier. */ + String newUuid(); +} diff --git a/core/src/main/java/com/google/adk/platform/package-info.java b/core/src/main/java/com/google/adk/platform/package-info.java new file mode 100644 index 000000000..c7eb66351 --- /dev/null +++ b/core/src/main/java/com/google/adk/platform/package-info.java @@ -0,0 +1,27 @@ +/* + * 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. + */ + +/** + * Seams for overriding system operations such as reading the current time and generating unique + * IDs. + * + *

By default ADK uses the wall clock and random UUIDs. Integrations that need customized + * timestamps and identifiers can install custom {@link com.google.adk.platform.TimeProvider} and + * {@link com.google.adk.platform.UuidProvider} implementations. Providers are carried on an + * invocation rather than in static state, which keeps them isolated to a single run and safe for + * concurrent invocations with independent providers. + */ +package com.google.adk.platform; diff --git a/core/src/main/java/com/google/adk/runner/Runner.java b/core/src/main/java/com/google/adk/runner/Runner.java index e13ce1235..d80a0a7a6 100644 --- a/core/src/main/java/com/google/adk/runner/Runner.java +++ b/core/src/main/java/com/google/adk/runner/Runner.java @@ -33,6 +33,8 @@ import com.google.adk.flows.llmflows.PersistBarrier; import com.google.adk.memory.BaseMemoryService; import com.google.adk.models.Model; +import com.google.adk.platform.TimeProvider; +import com.google.adk.platform.UuidProvider; import com.google.adk.plugins.Plugin; import com.google.adk.plugins.PluginManager; import com.google.adk.sessions.BaseSessionService; @@ -85,6 +87,8 @@ public class Runner { private final PluginManager pluginManager; @Nullable private final EventsCompactionConfig eventsCompactionConfig; @Nullable private final ContextCacheConfig contextCacheConfig; + private final TimeProvider timeProvider; + private final UuidProvider uuidProvider; private final ConcurrentMap activeSessionCompletables = new MapMaker().weakValues().makeMap(); @@ -97,6 +101,8 @@ public static class Builder { private BaseSessionService sessionService = new InMemorySessionService(); @Nullable private BaseMemoryService memoryService = null; private List plugins = ImmutableList.of(); + private TimeProvider timeProvider = TimeProvider.SYSTEM; + private UuidProvider uuidProvider = UuidProvider.SYSTEM; @CanIgnoreReturnValue public Builder app(App app) { @@ -151,6 +157,18 @@ public Builder plugins(Plugin... plugins) { return this; } + @CanIgnoreReturnValue + public Builder timeProvider(TimeProvider timeProvider) { + this.timeProvider = timeProvider; + return this; + } + + @CanIgnoreReturnValue + public Builder uuidProvider(UuidProvider uuidProvider) { + this.uuidProvider = uuidProvider; + return this; + } + public Runner build() { BaseAgent buildAgent; String buildAppName; @@ -198,7 +216,9 @@ public Runner build() { memoryService, buildPlugins, buildEventsCompactionConfig, - buildContextCacheConfig); + buildContextCacheConfig, + timeProvider, + uuidProvider); } } @@ -252,6 +272,30 @@ protected Runner( List plugins, @Nullable EventsCompactionConfig eventsCompactionConfig, @Nullable ContextCacheConfig contextCacheConfig) { + this( + agent, + appName, + artifactService, + sessionService, + memoryService, + plugins, + eventsCompactionConfig, + contextCacheConfig, + TimeProvider.SYSTEM, + UuidProvider.SYSTEM); + } + + private Runner( + BaseAgent agent, + String appName, + BaseArtifactService artifactService, + BaseSessionService sessionService, + @Nullable BaseMemoryService memoryService, + List plugins, + @Nullable EventsCompactionConfig eventsCompactionConfig, + @Nullable ContextCacheConfig contextCacheConfig, + TimeProvider timeProvider, + UuidProvider uuidProvider) { this.agent = agent; this.appName = appName; this.artifactService = artifactService; @@ -260,6 +304,8 @@ protected Runner( this.pluginManager = new PluginManager(plugins); this.eventsCompactionConfig = createEventsCompactionConfig(agent, eventsCompactionConfig); this.contextCacheConfig = contextCacheConfig; + this.timeProvider = timeProvider; + this.uuidProvider = uuidProvider; } /** @@ -352,7 +398,8 @@ private Single appendNewMessageToSession( // Appends only. We do not yield the event because it's not from the model. Event.Builder eventBuilder = Event.builder() - .id(Event.generateEventId()) + .id(invocationContext.newUuid()) + .timestamp(invocationContext.now().toEpochMilli()) .invocationId(invocationContext.invocationId()) .author("user") .content(newMessage); @@ -489,7 +536,7 @@ protected Flowable runAsyncImpl( return Flowable.defer( () -> { BaseAgent rootAgent = this.agent; - String invocationId = InvocationContext.newInvocationContextId(); + String invocationId = InvocationContext.newInvocationContextId(this.uuidProvider); // Pre-merge stateDelta so onUserMessageCallback can access it. // Safe: session is a copy; persistence still happens via appendNewMessageToSession. @@ -570,7 +617,8 @@ private Flowable runAgentWithUpdatedSession( .map( content -> Event.builder() - .id(Event.generateEventId()) + .id(contextWithUpdatedSession.newUuid()) + .timestamp(contextWithUpdatedSession.now().toEpochMilli()) .invocationId(contextWithUpdatedSession.invocationId()) .author("model") .content(content) @@ -678,6 +726,8 @@ private InvocationContext.Builder newInvocationContextBuilder(Session session) { .session(session) .eventsCompactionConfig(this.eventsCompactionConfig) .contextCacheConfig(this.contextCacheConfig) + .timeProvider(this.timeProvider) + .uuidProvider(this.uuidProvider) .agent(this.findAgentToRun(session, rootAgent)); } diff --git a/core/src/main/java/com/google/adk/sessions/InMemorySessionService.java b/core/src/main/java/com/google/adk/sessions/InMemorySessionService.java index 72a14cc4d..a3c9fecfa 100644 --- a/core/src/main/java/com/google/adk/sessions/InMemorySessionService.java +++ b/core/src/main/java/com/google/adk/sessions/InMemorySessionService.java @@ -20,6 +20,8 @@ import com.google.adk.events.Event; import com.google.adk.events.EventActions; +import com.google.adk.platform.TimeProvider; +import com.google.adk.platform.UuidProvider; import com.google.common.collect.ImmutableList; import com.google.errorprone.annotations.CanIgnoreReturnValue; import io.reactivex.rxjava3.core.Completable; @@ -32,7 +34,6 @@ import java.util.Map; import java.util.Objects; import java.util.Optional; -import java.util.UUID; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; import org.jspecify.annotations.Nullable; @@ -58,11 +59,21 @@ public final class InMemorySessionService implements BaseSessionService { // Structure: appName -> stateKey -> stateValue private final ConcurrentMap> appState; + private final TimeProvider timeProvider; + private final UuidProvider uuidProvider; + /** Creates a new instance of the in-memory session service with empty storage. */ public InMemorySessionService() { + this(TimeProvider.SYSTEM, UuidProvider.SYSTEM); + } + + /** Creates a new instance backed by the given time and UUID providers. */ + public InMemorySessionService(TimeProvider timeProvider, UuidProvider uuidProvider) { this.sessions = new ConcurrentHashMap<>(); this.userState = new ConcurrentHashMap<>(); this.appState = new ConcurrentHashMap<>(); + this.timeProvider = timeProvider; + this.uuidProvider = uuidProvider; } @Override @@ -87,7 +98,7 @@ public Single createSession( Optional.ofNullable(sessionId) .map(String::trim) .filter(s -> !s.isEmpty()) - .orElseGet(() -> UUID.randomUUID().toString()); + .orElseGet(uuidProvider::newUuid); // Ensure state map and events list are mutable for the new session ConcurrentMap initialState = @@ -99,7 +110,7 @@ public Single createSession( .appName(appName) .userId(userId) .state(initialState) - .lastUpdateTime(Instant.now()) + .lastUpdateTime(timeProvider.now()) .build(); sessions diff --git a/core/src/test/java/com/google/adk/agents/InvocationContextTest.java b/core/src/test/java/com/google/adk/agents/InvocationContextTest.java index e588a38ca..9609666b4 100644 --- a/core/src/test/java/com/google/adk/agents/InvocationContextTest.java +++ b/core/src/test/java/com/google/adk/agents/InvocationContextTest.java @@ -23,12 +23,15 @@ import com.google.adk.artifacts.BaseArtifactService; import com.google.adk.memory.BaseMemoryService; import com.google.adk.models.LlmCallsLimitExceededException; +import com.google.adk.platform.TimeProvider; +import com.google.adk.platform.UuidProvider; import com.google.adk.plugins.PluginManager; import com.google.adk.sessions.BaseSessionService; import com.google.adk.sessions.Session; import com.google.adk.summarizer.EventsCompactionConfig; import com.google.common.collect.ImmutableMap; import com.google.genai.types.Content; +import java.time.Instant; import java.util.HashMap; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; @@ -258,6 +261,77 @@ public void testNewInvocationContextId() { .matches("^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$"); } + @Test + public void testNewInvocationContextId_withCustomProvider() { + UuidProvider provider = () -> "deterministic-uuid"; + + String id = InvocationContext.newInvocationContextId(provider); + + assertThat(id).isEqualTo("e-deterministic-uuid"); + } + + @Test + public void providers_defaultToSystem() { + InvocationContext context = + InvocationContext.builder() + .sessionService(mockSessionService) + .artifactService(mockArtifactService) + .agent(mockAgent) + .session(session) + .build(); + + assertThat(context.timeProvider()).isEqualTo(TimeProvider.SYSTEM); + assertThat(context.uuidProvider()).isEqualTo(UuidProvider.SYSTEM); + assertThat(context.newUuid()) + .matches("^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$"); + Instant before = Instant.now(); + Instant now = context.now(); + assertThat(now).isAtLeast(before); + } + + @Test + public void providers_customAreCarriedAndUsedByAccessors() { + Instant fixed = Instant.ofEpochMilli(1234L); + TimeProvider fixedClock = () -> fixed; + UuidProvider fixedUuids = () -> "fixed-uuid"; + + InvocationContext context = + InvocationContext.builder() + .sessionService(mockSessionService) + .artifactService(mockArtifactService) + .agent(mockAgent) + .session(session) + .timeProvider(fixedClock) + .uuidProvider(fixedUuids) + .build(); + + assertThat(context.timeProvider()).isEqualTo(fixedClock); + assertThat(context.uuidProvider()).isEqualTo(fixedUuids); + assertThat(context.now()).isEqualTo(fixed); + assertThat(context.newUuid()).isEqualTo("fixed-uuid"); + } + + @Test + public void toBuilder_carriesProviders() { + TimeProvider fixedClock = () -> Instant.ofEpochMilli(1234L); + UuidProvider fixedUuids = () -> "fixed-uuid"; + + InvocationContext originalContext = + InvocationContext.builder() + .sessionService(mockSessionService) + .artifactService(mockArtifactService) + .agent(mockAgent) + .session(session) + .timeProvider(fixedClock) + .uuidProvider(fixedUuids) + .build(); + + InvocationContext copiedContext = originalContext.toBuilder().build(); + + assertThat(copiedContext.timeProvider()).isEqualTo(fixedClock); + assertThat(copiedContext.uuidProvider()).isEqualTo(fixedUuids); + } + @Test public void testEquals_sameObject() { InvocationContext context = diff --git a/core/src/test/java/com/google/adk/flows/llmflows/FunctionsTest.java b/core/src/test/java/com/google/adk/flows/llmflows/FunctionsTest.java index b2ffc4443..42cbb98da 100644 --- a/core/src/test/java/com/google/adk/flows/llmflows/FunctionsTest.java +++ b/core/src/test/java/com/google/adk/flows/llmflows/FunctionsTest.java @@ -25,6 +25,7 @@ import com.google.adk.agents.RunConfig; import com.google.adk.agents.RunConfig.ToolExecutionMode; import com.google.adk.events.Event; +import com.google.adk.platform.UuidProvider; import com.google.adk.testing.TestUtils; import com.google.adk.tools.BaseTool; import com.google.adk.tools.ToolContext; @@ -76,6 +77,31 @@ public final class FunctionsTest { .content(Content.fromParts(Part.fromFunctionCall("other_function", ImmutableMap.of()))) .build(); + @Test + public void generateClientFunctionCallId_withProvider_derivesFromProvider() { + UuidProvider provider = () -> "deterministic-uuid"; + + String id = Functions.generateClientFunctionCallId(provider); + + assertThat(id).isEqualTo("adk-deterministic-uuid"); + } + + @Test + public void populateClientFunctionCallId_withProvider_usesProvider() { + Event event = + Event.builder() + .id("event1") + .invocationId("invocation1") + .author("agent") + .content(Content.fromParts(Part.fromFunctionCall("some_function", ImmutableMap.of()))) + .build(); + + Functions.populateClientFunctionCallId(event, () -> "deterministic-uuid"); + + Part populatedPart = event.content().get().parts().get().get(0); + assertThat(populatedPart.functionCall().get().id()).hasValue("adk-deterministic-uuid"); + } + @Test public void handleFunctionCalls_noFunctionCalls() { InvocationContext invocationContext = createInvocationContext(createRootAgent()); diff --git a/core/src/test/java/com/google/adk/platform/TimeProviderTest.java b/core/src/test/java/com/google/adk/platform/TimeProviderTest.java new file mode 100644 index 000000000..51a0293ad --- /dev/null +++ b/core/src/test/java/com/google/adk/platform/TimeProviderTest.java @@ -0,0 +1,47 @@ +/* + * 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.platform; + +import static com.google.common.truth.Truth.assertThat; + +import java.time.Instant; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public final class TimeProviderTest { + + @Test + public void system_returnsWallClockTime() { + Instant before = Instant.now(); + Instant now = TimeProvider.SYSTEM.now(); + Instant after = Instant.now(); + + assertThat(now).isAtLeast(before); + assertThat(now).isAtMost(after); + } + + @Test + public void customProvider_returnsFixedTime() { + Instant fixed = Instant.ofEpochMilli(1234L); + TimeProvider provider = () -> fixed; + + assertThat(provider.now()).isEqualTo(fixed); + assertThat(provider.now()).isEqualTo(fixed); + } +} diff --git a/core/src/test/java/com/google/adk/platform/UuidProviderTest.java b/core/src/test/java/com/google/adk/platform/UuidProviderTest.java new file mode 100644 index 000000000..4248d01c1 --- /dev/null +++ b/core/src/test/java/com/google/adk/platform/UuidProviderTest.java @@ -0,0 +1,55 @@ +/* + * 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.platform; + +import static com.google.common.truth.Truth.assertThat; + +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public final class UuidProviderTest { + + @Test + public void system_returnsUniqueUuidStrings() { + String first = UuidProvider.SYSTEM.newUuid(); + String second = UuidProvider.SYSTEM.newUuid(); + + assertThat(first).isNotEmpty(); + assertThat(first).isNotEqualTo(second); + assertThat(first).matches("^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$"); + } + + @Test + public void customProvider_returnsDeterministicSequence() { + UuidProvider provider = + new UuidProvider() { + private final AtomicInteger counter = new AtomicInteger(); + + @Override + public String newUuid() { + return String.format("uuid-%04d", counter.getAndIncrement()); + } + }; + + assertThat(provider.newUuid()).isEqualTo("uuid-0000"); + assertThat(provider.newUuid()).isEqualTo("uuid-0001"); + assertThat(provider.newUuid()).isEqualTo("uuid-0002"); + } +} diff --git a/core/src/test/java/com/google/adk/runner/RunnerTest.java b/core/src/test/java/com/google/adk/runner/RunnerTest.java index 9d70cf4a0..33bc4585c 100644 --- a/core/src/test/java/com/google/adk/runner/RunnerTest.java +++ b/core/src/test/java/com/google/adk/runner/RunnerTest.java @@ -51,6 +51,8 @@ import com.google.adk.flows.llmflows.Functions; import com.google.adk.models.LlmRequest; import com.google.adk.models.LlmResponse; +import com.google.adk.platform.TimeProvider; +import com.google.adk.platform.UuidProvider; import com.google.adk.plugins.BasePlugin; import com.google.adk.sessions.BaseSessionService; import com.google.adk.sessions.GetSessionConfig; @@ -1531,6 +1533,68 @@ public void runAsync_ensureEventsAreAppendedInOrder() throws Exception { .inOrder(); } + /** + * Verification of custom time/UUID provider behavior: running identical input twice with the + * same deterministic clock and UUID provider yields byte-identical event ids, timestamps, and + * invocation ids. + */ + @Test + public void runAsync_withDeterministicProviders_replayProducesIdenticalEvents() { + List firstRun = runWithDeterministicProviders(); + List secondRun = runWithDeterministicProviders(); + + assertThat(firstRun).isNotEmpty(); + + List firstIds = firstRun.stream().map(Event::id).toList(); + List secondIds = secondRun.stream().map(Event::id).toList(); + assertThat(firstIds).isEqualTo(secondIds); + + List firstTimestamps = firstRun.stream().map(Event::timestamp).toList(); + List secondTimestamps = secondRun.stream().map(Event::timestamp).toList(); + assertThat(firstTimestamps).isEqualTo(secondTimestamps); + + List firstInvocationIds = firstRun.stream().map(Event::invocationId).toList(); + List secondInvocationIds = secondRun.stream().map(Event::invocationId).toList(); + assertThat(firstInvocationIds).isEqualTo(secondInvocationIds); + + // The injected clock seam is actually used: every event carries the fixed timestamp. + assertThat(firstTimestamps).doesNotContain(null); + assertThat(firstTimestamps.stream().distinct().toList()).containsExactly(1234L); + // The invocation id is minted from the injected UUID provider. + assertThat(firstInvocationIds.stream().distinct().toList()).containsExactly("e-uuid-0000"); + } + + private List runWithDeterministicProviders() { + TimeProvider fixedClock = () -> Instant.ofEpochMilli(1234L); + UuidProvider sequentialUuids = + new UuidProvider() { + private final AtomicInteger counter = new AtomicInteger(); + + @Override + public String newUuid() { + return String.format("uuid-%04d", counter.getAndIncrement()); + } + }; + LlmAgent deterministicAgent = + createTestAgentBuilder(createTestLlm(createLlmResponse(createContent("from llm")))).build(); + Runner deterministicRunner = + Runner.builder() + .app(App.builder().name("test").rootAgent(deterministicAgent).build()) + .sessionService(new InMemorySessionService(fixedClock, sequentialUuids)) + .timeProvider(fixedClock) + .uuidProvider(sequentialUuids) + .build(); + Session deterministicSession = + deterministicRunner + .sessionService() + .createSession("test", "user", new ConcurrentHashMap(), "fixed-session") + .blockingGet(); + return deterministicRunner + .runAsync("user", deterministicSession.id(), createContent("hi")) + .toList() + .blockingGet(); + } + private Content createContent(String text) { return Content.builder().parts(Part.builder().text(text).build()).build(); } diff --git a/core/src/test/java/com/google/adk/sessions/InMemorySessionServiceTest.java b/core/src/test/java/com/google/adk/sessions/InMemorySessionServiceTest.java index 6a271efac..9b516af01 100644 --- a/core/src/test/java/com/google/adk/sessions/InMemorySessionServiceTest.java +++ b/core/src/test/java/com/google/adk/sessions/InMemorySessionServiceTest.java @@ -19,6 +19,8 @@ import com.google.adk.events.Event; import com.google.adk.events.EventActions; +import com.google.adk.platform.TimeProvider; +import com.google.adk.platform.UuidProvider; import io.reactivex.rxjava3.core.Single; import java.lang.reflect.Field; import java.time.Instant; @@ -66,6 +68,18 @@ public void lifecycle_createSession() { assertThat(session.state()).isEmpty(); } + @Test + public void createSession_withDeterministicProviders_usesProvidersForIdAndTime() { + TimeProvider fixedClock = () -> Instant.ofEpochMilli(1234L); + UuidProvider fixedUuids = () -> "fixed-session-id"; + InMemorySessionService sessionService = new InMemorySessionService(fixedClock, fixedUuids); + + Session session = sessionService.createSession("app-name", "user-id").blockingGet(); + + assertThat(session.id()).isEqualTo("fixed-session-id"); + assertThat(session.lastUpdateTime()).isEqualTo(Instant.ofEpochMilli(1234L)); + } + @Test public void lifecycle_getSession() { InMemorySessionService sessionService = new InMemorySessionService();