From caccc7cbbc3198b35bd769d20d72c1cd68a08d4c Mon Sep 17 00:00:00 2001 From: DABH Date: Tue, 9 Jun 2026 14:36:19 -0500 Subject: [PATCH] feat: add injectable time and UUID providers for timestamps and ids ADK generates timestamps and ids by calling Instant.now() and UUID.randomUUID() directly, leaving callers no way to control them. This blocks integrations that need to supply their own, such as durable-execution hosts that replay an invocation deterministically. adk-python (platform.time / platform.uuid) and adk-go (platform.Now / platform.NewUUID) already expose an equivalent seam. Give every invocation a java.time.InstantSource and a UuidProvider, a new functional interface in the leaf com.google.adk.platform package whose SYSTEM default is backed by UUID.randomUUID(). Both are configured once on Runner.Builder (instantSource(), uuidProvider()) and threaded as data through InvocationContext, which exposes them as instantSource(), uuidProvider(), now(), and newUuid(). Threading them as data rather than through an ambient ThreadLocal keeps them visible on whatever RxJava worker thread builds an event. The Runner passes them to the default InMemorySessionService it constructs, to the LlmEventSummarizer it builds for event compaction, and to every InvocationContext it creates for new, resumed, and live invocations. AgentTool's nested runner inherits the parent invocation's providers. Defaults keep their per-builder sharing semantics: Runner.Builder creates its default session service once and reuses it for every build(), an explicit sessionService(null) is still rejected at build() rather than in the setter, and InvocationContext.Builder mints its default invocation id once and shares it across repeated build() calls. Every generated id and timestamp on the Runner -> agent -> flow -> session-event path now derives from the invocation's providers: the invocation id; user, model, tool, callback, code-execution, and output-schema events; the resumability checkpoint and state-delta events; replayed function-call re-ids; the client function-call ids the flow mints (including the tool-confirmation request); compaction events; remote A2A agent events and outbound message ids; the PlannerAgent result event; and the session id and lastUpdateTime of the in-memory and Firestore session services. Event.generateEventId() and the Event.Builder timestamp default keep reading the system sources, so the public Event API is unchanged; their Javadoc points custom agents, tools, and plugins at InvocationContext.newUuid() and InvocationContext.now(). Event.toBuilder() now copies the timestamp unconditionally. It used to drop a zero timestamp so that BaseLlmFlow's pre-call event template, which zeroed its own, would be re-stamped from the wall clock on rebuild; with a clock fixed at the epoch that guard silently replaced the timestamp of every rebuilt event. The template is now stamped from the invocation's clock, and model-response events are timestamped when the response arrives, as in adk-python. Code-execution and tool-confirmation events, which had no id before, now carry one. BaseLlmFlow no longer re-ids its per-step template after the model call: nothing read it by then. Runner.Builder.instantSource(null) and uuidProvider(null), and the matching InvocationContext.Builder setters, select the system defaults, like scheduler(null) does, so a nested run started from a mocked or partially configured context still works. AgentTool passes the caller's scheduler and providers to the runner it nests. Out of scope: function-call ids a model adapter assigns itself while aggregating a streamed response, the A2A server executor's protocol ids, the Firestore create token, the wall-clock timestamp CodeExecutorContext writes into session state, and telemetry and analytics timestamps and ids. PlannerAgentTest is converted to JUnit 5 so the Jupiter-only planners module actually runs it. --- .../google/adk/a2a/agent/RemoteA2AAgent.java | 17 +- .../adk/a2a/converters/EventConverter.java | 42 +- .../adk/a2a/converters/ResponseConverter.java | 10 +- .../adk/a2a/agent/RemoteA2AAgentTest.java | 35 ++ .../a2a/converters/EventConverterTest.java | 10 + .../a2a/converters/ResponseConverterTest.java | 36 ++ .../adk/sessions/FirestoreSessionService.java | 25 +- .../sessions/FirestoreSessionServiceTest.java | 27 ++ .../com/google/adk/agents/PlannerAgent.java | 3 +- .../google/adk/agents/PlannerAgentTest.java | 43 +- .../java/com/google/adk/agents/BaseAgent.java | 9 +- .../google/adk/agents/InvocationContext.java | 97 +++- .../java/com/google/adk/events/Event.java | 25 +- .../adk/flows/llmflows/BaseLlmFlow.java | 30 +- .../adk/flows/llmflows/CodeExecution.java | 6 + .../google/adk/flows/llmflows/Functions.java | 33 +- .../adk/flows/llmflows/OutputSchema.java | 3 +- .../google/adk/models/FunctionCallIds.java | 12 +- .../com/google/adk/platform/UuidProvider.java | 34 ++ .../com/google/adk/platform/package-info.java | 26 ++ .../java/com/google/adk/runner/Runner.java | 87 +++- .../adk/sessions/InMemorySessionService.java | 26 +- .../adk/summarizer/LlmEventSummarizer.java | 32 +- .../java/com/google/adk/tools/AgentTool.java | 14 +- .../adk/agents/InvocationContextTest.java | 182 ++++++++ .../java/com/google/adk/events/EventTest.java | 15 + .../adk/flows/llmflows/CodeExecutionTest.java | 25 +- .../adk/flows/llmflows/FunctionsTest.java | 102 ++++- .../adk/flows/llmflows/OutputSchemaTest.java | 18 + .../adk/models/FunctionCallIdsTest.java | 11 + .../google/adk/platform/UuidProviderTest.java | 55 +++ .../adk/runner/RunnerResumabilityTest.java | 64 +++ .../com/google/adk/runner/RunnerTest.java | 413 ++++++++++++++++++ .../sessions/InMemorySessionServiceTest.java | 17 + .../summarizer/LlmEventSummarizerTest.java | 34 ++ .../com/google/adk/tools/AgentToolTest.java | 75 ++++ 36 files changed, 1556 insertions(+), 137 deletions(-) create mode 100644 core/src/main/java/com/google/adk/platform/UuidProvider.java create mode 100644 core/src/main/java/com/google/adk/platform/package-info.java create mode 100644 core/src/test/java/com/google/adk/platform/UuidProviderTest.java diff --git a/a2a/src/main/java/com/google/adk/a2a/agent/RemoteA2AAgent.java b/a2a/src/main/java/com/google/adk/a2a/agent/RemoteA2AAgent.java index 571493b8d..bcad60b15 100644 --- a/a2a/src/main/java/com/google/adk/a2a/agent/RemoteA2AAgent.java +++ b/a2a/src/main/java/com/google/adk/a2a/agent/RemoteA2AAgent.java @@ -48,12 +48,10 @@ import io.reactivex.rxjava3.core.BackpressureStrategy; import io.reactivex.rxjava3.core.Flowable; import io.reactivex.rxjava3.core.FlowableEmitter; -import java.time.Instant; import java.util.ArrayList; import java.util.List; import java.util.Objects; import java.util.Optional; -import java.util.UUID; import java.util.function.BiConsumer; import org.jspecify.annotations.Nullable; import org.slf4j.Logger; @@ -193,8 +191,9 @@ public boolean isStreaming() { return streaming; } - private Message.Builder newA2AMessage(Message.Role role, List> parts) { - return new Message.Builder().messageId(UUID.randomUUID().toString()).role(role).parts(parts); + private Message.Builder newA2AMessage( + InvocationContext invocationContext, Message.Role role, List> parts) { + return new Message.Builder().messageId(invocationContext.newUuid()).role(role).parts(parts); } private Message prepareMessage(InvocationContext invocationContext) { @@ -202,13 +201,15 @@ private Message prepareMessage(InvocationContext invocationContext) { if (userCall != null) { ImmutableList> parts = EventConverter.contentToParts(userCall.content(), userCall.partial().orElse(false)); - return newA2AMessage(Message.Role.USER, parts) + return newA2AMessage(invocationContext, Message.Role.USER, parts) .taskId(EventConverter.taskId(userCall)) .contextId(EventConverter.contextId(userCall)) .build(); } return newA2AMessage( - Message.Role.USER, EventConverter.messagePartsFromContext(invocationContext)) + invocationContext, + Message.Role.USER, + EventConverter.messagePartsFromContext(invocationContext)) .build(); } @@ -568,11 +569,11 @@ private Event createAggregatedEvent(Content content, @Nullable ClientEvent trigg } return Event.builder() - .id(UUID.randomUUID().toString()) + .id(invocationContext.newUuid()) .invocationId(invocationContext.invocationId()) .author(agentName) .content(content) - .timestamp(Instant.now().toEpochMilli()) + .timestamp(invocationContext.now().toEpochMilli()) .customMetadata(aggMetadataBuilder.build()) .build(); } diff --git a/a2a/src/main/java/com/google/adk/a2a/converters/EventConverter.java b/a2a/src/main/java/com/google/adk/a2a/converters/EventConverter.java index 71573070e..df10a9c85 100644 --- a/a2a/src/main/java/com/google/adk/a2a/converters/EventConverter.java +++ b/a2a/src/main/java/com/google/adk/a2a/converters/EventConverter.java @@ -27,7 +27,6 @@ import java.util.Collection; import java.util.List; import java.util.Optional; -import java.util.UUID; import org.jspecify.annotations.Nullable; /** Converter for ADK Events to A2A Messages. */ @@ -144,29 +143,36 @@ public static ImmutableList> messagePartsFromContext(InvocationContext c } List events = context.session().events(); int lastResponseIndex = -1; - String contextId = ""; for (int i = events.size() - 1; i >= 0; i--) { Event event = events.get(i); if (event.author().equals(context.agent().name())) { lastResponseIndex = i; - contextId = contextId(event); break; } } ImmutableList.Builder> partsBuilder = ImmutableList.builder(); for (int i = lastResponseIndex + 1; i < events.size(); i++) { Event event = events.get(i); + Optional content; + boolean partial; if (!event.author().equals("user") && !event.author().equals(context.agent().name())) { - event = presentAsUserMessage(event, contextId); + // Another agent's output is rephrased as a complete user message. + content = presentAsUserMessage(event); + partial = false; + } else { + content = event.content(); + partial = event.partial().orElse(false); } - contentToParts(event.content(), event.partial().orElse(false)).forEach(partsBuilder::add); + contentToParts(content, partial).forEach(partsBuilder::add); } return partsBuilder.build(); } - private static Event presentAsUserMessage(Event event, String contextId) { - Event.Builder userEvent = - new Event.Builder().id(UUID.randomUUID().toString()).invocationId(contextId).author("user"); + /** + * Rephrases the content of an event authored by another agent as a user message, or returns empty + * when the event has no parts to present. + */ + private static Optional presentAsUserMessage(Event event) { ImmutableList parts = event.content().flatMap(Content::parts).stream() .flatMap(Collection::stream) @@ -176,20 +182,18 @@ private static Event presentAsUserMessage(Event event, String contextId) { .map(part -> PartConverter.remoteCallAsUserPart(event.author(), part)) .collect(toImmutableList()); if (parts.isEmpty()) { - return userEvent.build(); + return Optional.empty(); } com.google.genai.types.Part forContext = com.google.genai.types.Part.builder().text("For context:").build(); - return userEvent - .content( - Content.builder() - .parts( - ImmutableList.builder() - .add(forContext) - .addAll(parts) - .build()) - .build()) - .build(); + return Optional.of( + Content.builder() + .parts( + ImmutableList.builder() + .add(forContext) + .addAll(parts) + .build()) + .build()); } private static String metadataValue(Event event, String key) { diff --git a/a2a/src/main/java/com/google/adk/a2a/converters/ResponseConverter.java b/a2a/src/main/java/com/google/adk/a2a/converters/ResponseConverter.java index cebdc4b52..9e1a8687d 100644 --- a/a2a/src/main/java/com/google/adk/a2a/converters/ResponseConverter.java +++ b/a2a/src/main/java/com/google/adk/a2a/converters/ResponseConverter.java @@ -44,12 +44,10 @@ import io.a2a.spec.TaskArtifactUpdateEvent; import io.a2a.spec.TaskState; import io.a2a.spec.TaskStatusUpdateEvent; -import java.time.Instant; import java.util.List; import java.util.Map; import java.util.Objects; import java.util.Optional; -import java.util.UUID; import org.jspecify.annotations.Nullable; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -402,12 +400,12 @@ private static void logDroppedMetadata(A2AMetadataKey key, Exception e) { private static Event emptyEvent(InvocationContext invocationContext) { Event.Builder builder = Event.builder() - .id(UUID.randomUUID().toString()) + .id(invocationContext.newUuid()) .invocationId(invocationContext.invocationId()) .author(invocationContext.agent().name()) .branch(invocationContext.branch().orElse(null)) .content(Content.builder().role("user").parts(ImmutableList.of()).build()) - .timestamp(Instant.now().toEpochMilli()); + .timestamp(invocationContext.now().toEpochMilli()); return builder.build(); } @@ -421,10 +419,10 @@ private static boolean hasParts(Event event) { private static Event.Builder remoteAgentEventBuilder(InvocationContext invocationContext) { return Event.builder() - .id(UUID.randomUUID().toString()) + .id(invocationContext.newUuid()) .invocationId(invocationContext.invocationId()) .author(invocationContext.agent().name()) .branch(invocationContext.branch().orElse(null)) - .timestamp(Instant.now().toEpochMilli()); + .timestamp(invocationContext.now().toEpochMilli()); } } diff --git a/a2a/src/test/java/com/google/adk/a2a/agent/RemoteA2AAgentTest.java b/a2a/src/test/java/com/google/adk/a2a/agent/RemoteA2AAgentTest.java index 66a1f0cde..2bb450003 100644 --- a/a2a/src/test/java/com/google/adk/a2a/agent/RemoteA2AAgentTest.java +++ b/a2a/src/test/java/com/google/adk/a2a/agent/RemoteA2AAgentTest.java @@ -67,6 +67,7 @@ import io.reactivex.rxjava3.core.Maybe; import io.reactivex.rxjava3.plugins.RxJavaPlugins; import io.reactivex.rxjava3.subscribers.TestSubscriber; +import java.time.Instant; import java.util.AbstractList; import java.util.ArrayList; import java.util.Collections; @@ -76,6 +77,7 @@ import java.util.Objects; import java.util.concurrent.CountDownLatch; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; import java.util.function.BiConsumer; import java.util.function.Consumer; @@ -220,6 +222,39 @@ public void runAsync_aggregatesPartialEvents() { assertResponseMetadata(finalEvent); } + @Test + @SuppressWarnings("unchecked") // cast for Mockito + public void runAsync_usesInvocationProvidersForMessageIdAndEvents() { + RemoteA2AAgent agent = createAgent(); + AtomicInteger counter = new AtomicInteger(); + InvocationContext providerContext = + invocationContext.toBuilder() + .instantSource(() -> Instant.ofEpochMilli(1234L)) + .uuidProvider(() -> String.format("uuid-%04d", counter.getAndIncrement())) + .build(); + mockStreamResponse( + consumer -> { + consumer.accept(createPartialEvent("Hello ", true, false), agentCard); + consumer.accept(createPartialEvent("World!", true, false), agentCard); + consumer.accept(createFinalEvent("Final artifact content"), agentCard); + }); + + List events = agent.runAsync(providerContext).toList().blockingGet(); + + ArgumentCaptor messageCaptor = ArgumentCaptor.forClass(Message.class); + verify(mockClient) + .sendMessage(messageCaptor.capture(), any(List.class), any(Consumer.class), any()); + // The outbound A2A message id is the first id drawn from the invocation's provider: building + // the request parts draws none. + assertThat(messageCaptor.getValue().getMessageId()).isEqualTo("uuid-0000"); + assertThat(events).hasSize(4); + assertThat(events.stream().map(Event::id).distinct()).hasSize(4); + for (Event event : events) { + assertThat(event.id()).startsWith("uuid-"); + assertThat(event.timestamp()).isEqualTo(1234L); + } + } + @Test public void runAsync_aggregatesInterleavedFunctionCalls() { RemoteA2AAgent agent = createAgent(); diff --git a/a2a/src/test/java/com/google/adk/a2a/converters/EventConverterTest.java b/a2a/src/test/java/com/google/adk/a2a/converters/EventConverterTest.java index 74292618d..51fc3184f 100644 --- a/a2a/src/test/java/com/google/adk/a2a/converters/EventConverterTest.java +++ b/a2a/src/test/java/com/google/adk/a2a/converters/EventConverterTest.java @@ -196,6 +196,7 @@ public void testMessagePartsFromContext() { .build(), Event.builder() .author("other_agent") + .partial(true) .content( Content.builder() .parts(ImmutableList.of(Part.builder().text("hey").build())) @@ -208,12 +209,21 @@ public void testMessagePartsFromContext() { .session(session) .sessionService(new InMemorySessionService()) .agent(agent) + // An explicit invocation id keeps build() from drawing one; rephrasing another agent's + // output must draw none. + .invocationId("invocation-id") + .uuidProvider( + () -> { + throw new AssertionError("messagePartsFromContext must not draw ids"); + }) .build(); ImmutableList> parts = EventConverter.messagePartsFromContext(ctx); assertThat(parts).hasSize(2); assertThat(((TextPart) parts.get(0)).getText()).isEqualTo("For context:"); assertThat(((TextPart) parts.get(1)).getText()).isEqualTo("[other_agent] said: hey"); + // Another agent's output is rephrased as a complete message even when its event was partial. + assertThat(parts.get(1).getMetadata()).doesNotContainKey(A2AMetadataKey.PARTIAL.getType()); } private static final class TestAgent extends BaseAgent { diff --git a/a2a/src/test/java/com/google/adk/a2a/converters/ResponseConverterTest.java b/a2a/src/test/java/com/google/adk/a2a/converters/ResponseConverterTest.java index 684009df9..b33b1ccf5 100644 --- a/a2a/src/test/java/com/google/adk/a2a/converters/ResponseConverterTest.java +++ b/a2a/src/test/java/com/google/adk/a2a/converters/ResponseConverterTest.java @@ -49,6 +49,7 @@ import io.a2a.spec.TaskStatusUpdateEvent; import io.a2a.spec.TextPart; import io.reactivex.rxjava3.core.Flowable; +import java.time.Instant; import java.util.Arrays; import java.util.List; import java.util.Optional; @@ -140,6 +141,41 @@ public void messageToEvent_convertsMessage() { assertThat(content.parts().get().get(0).text()).hasValue("test-message"); } + @Test + public void messageToEvent_usesInvocationProvidersForIdAndTimestamp() { + InvocationContext providerContext = + invocationContext.toBuilder() + .instantSource(() -> Instant.ofEpochMilli(1234L)) + .uuidProvider(() -> "fixed-uuid") + .build(); + Message a2aMessage = + new Message.Builder() + .messageId("msg-1") + .role(Message.Role.USER) + .parts(ImmutableList.of(new TextPart("test-message"))) + .build(); + + Event event = ResponseConverter.messageToEvent(a2aMessage, providerContext); + + assertThat(event.id()).isEqualTo("fixed-uuid"); + assertThat(event.timestamp()).isEqualTo(1234L); + } + + @Test + public void taskToEvent_withNoMessage_emptyEventUsesInvocationProviders() { + InvocationContext providerContext = + invocationContext.toBuilder() + .instantSource(() -> Instant.ofEpochMilli(1234L)) + .uuidProvider(() -> "fixed-uuid") + .build(); + Task task = testTask().status(new TaskStatus(TaskState.WORKING, null, null)).build(); + + Event event = ResponseConverter.taskToEvent(task, providerContext); + + assertThat(event.id()).isEqualTo("fixed-uuid"); + assertThat(event.timestamp()).isEqualTo(1234L); + } + @Test public void taskToEvent_withArtifacts_returnsEventFromLastArtifact() { io.a2a.spec.Part a2aPart = new TextPart("Artifact content"); diff --git a/contrib/firestore-session-service/src/main/java/com/google/adk/sessions/FirestoreSessionService.java b/contrib/firestore-session-service/src/main/java/com/google/adk/sessions/FirestoreSessionService.java index 621421b31..54f91a858 100644 --- a/contrib/firestore-session-service/src/main/java/com/google/adk/sessions/FirestoreSessionService.java +++ b/contrib/firestore-session-service/src/main/java/com/google/adk/sessions/FirestoreSessionService.java @@ -17,6 +17,7 @@ import com.google.adk.events.Event; import com.google.adk.events.EventActions; +import com.google.adk.platform.UuidProvider; import com.google.adk.utils.ApiFutureUtils; import com.google.adk.utils.Constants; import com.google.api.core.ApiFuture; @@ -38,6 +39,7 @@ import io.reactivex.rxjava3.core.Maybe; import io.reactivex.rxjava3.core.Single; import java.time.Instant; +import java.time.InstantSource; import java.util.ArrayList; import java.util.HashMap; import java.util.HashSet; @@ -64,6 +66,8 @@ public class FirestoreSessionService implements BaseSessionService { private static final Logger logger = LoggerFactory.getLogger(FirestoreSessionService.class); private final Firestore firestore; + private final InstantSource instantSource; + private final UuidProvider uuidProvider; private static final String ROOT_COLLECTION_NAME = Constants.ROOT_COLLECTION_NAME; private static final String EVENTS_SUBCOLLECTION_NAME = Constants.EVENTS_SUBCOLLECTION_NAME; private static final String APP_STATE_COLLECTION = Constants.APP_STATE_COLLECTION; @@ -81,7 +85,22 @@ public class FirestoreSessionService implements BaseSessionService { /** Constructor for FirestoreSessionService. */ public FirestoreSessionService(Firestore firestore) { - this.firestore = firestore; + this(firestore, InstantSource.system(), UuidProvider.SYSTEM); + } + + /** + * Creates a session service whose generated session IDs and {@code lastUpdateTime} values come + * from the given providers instead of {@link java.util.UUID#randomUUID()} and the wall clock. + * + * @param firestore The Firestore client. + * @param instantSource Supplies the {@code lastUpdateTime} of a newly created session. + * @param uuidProvider Supplies the session ID when {@code createSession} is not given one. + */ + public FirestoreSessionService( + Firestore firestore, InstantSource instantSource, UuidProvider uuidProvider) { + this.firestore = Objects.requireNonNull(firestore, "firestore cannot be null"); + this.instantSource = Objects.requireNonNull(instantSource, "instantSource cannot be null"); + this.uuidProvider = Objects.requireNonNull(uuidProvider, "uuidProvider cannot be null"); } /** Gets the sessions collection reference for a given userId. */ @@ -121,7 +140,7 @@ public Single createSession( Optional.ofNullable(sessionId) .map(String::trim) .filter(s -> !s.isEmpty()) - .orElseGet(() -> UUID.randomUUID().toString()); + .orElseGet(uuidProvider::newUuid); ConcurrentMap initialState = (state == null) ? new ConcurrentHashMap<>() : new ConcurrentHashMap<>(state); @@ -131,7 +150,7 @@ public Single createSession( resolvedSessionId, initialState); List initialEvents = new ArrayList<>(); - Instant now = Instant.now(); + Instant now = instantSource.instant(); Session newSession = Session.builder(resolvedSessionId) .appName(appName) diff --git a/contrib/firestore-session-service/src/test/java/com/google/adk/sessions/FirestoreSessionServiceTest.java b/contrib/firestore-session-service/src/test/java/com/google/adk/sessions/FirestoreSessionServiceTest.java index 9b7d1a522..ba95dc84c 100644 --- a/contrib/firestore-session-service/src/test/java/com/google/adk/sessions/FirestoreSessionServiceTest.java +++ b/contrib/firestore-session-service/src/test/java/com/google/adk/sessions/FirestoreSessionServiceTest.java @@ -318,6 +318,33 @@ void createSession_withSessionId_returnsNewSession() { verify(mockSessionDocRef).create(anyMap()); } + /** Tests that createSession draws the generated session ID and update time from the providers. */ + @Test + void createSession_withInjectedProviders_usesProvidersForIdAndUpdateTime() { + // Arrange + FirestoreSessionService providerService = + new FirestoreSessionService( + mockDb, () -> Instant.ofEpochMilli(1234L), () -> "fixed-session-id"); + when(mockSessionsCollection.document("fixed-session-id")).thenReturn(mockSessionDocRef); + + // Act + TestObserver testObserver = + providerService.createSession(APP_NAME, USER_ID, new ConcurrentHashMap<>(), null).test(); + + // Assert + testObserver.awaitCount(1); + testObserver.assertComplete(); + testObserver.assertValue( + session -> { + assertThat(session.id()).isEqualTo("fixed-session-id"); + assertThat(session.lastUpdateTime()).isEqualTo(Instant.ofEpochMilli(1234L)); + return true; + }); + verify(mockSessionDocRef).create(sessionDataCaptor.capture()); + assertThat(sessionDataCaptor.getValue()) + .containsEntry(Constants.KEY_UPDATE_TIME, Instant.ofEpochMilli(1234L).toString()); + } + /** Tests that createSession creates a new session with a generated session ID. */ @Test void createSession_withNullSessionId_generatesNewId() { diff --git a/contrib/planners/src/main/java/com/google/adk/agents/PlannerAgent.java b/contrib/planners/src/main/java/com/google/adk/agents/PlannerAgent.java index 909845b16..bd07b8ff5 100644 --- a/contrib/planners/src/main/java/com/google/adk/agents/PlannerAgent.java +++ b/contrib/planners/src/main/java/com/google/adk/agents/PlannerAgent.java @@ -132,7 +132,8 @@ private Flowable executeActionAndContinue( if (action instanceof PlannerAction.DoneWithResult doneWithResult) { Event resultEvent = Event.builder() - .id(Event.generateEventId()) + .id(invocationContext.newUuid()) + .timestamp(invocationContext.now().toEpochMilli()) .invocationId(invocationContext.invocationId()) .author(name()) .branch(invocationContext.branch().orElse(null)) diff --git a/contrib/planners/src/test/java/com/google/adk/agents/PlannerAgentTest.java b/contrib/planners/src/test/java/com/google/adk/agents/PlannerAgentTest.java index 9e16205b4..1031ef0ca 100644 --- a/contrib/planners/src/test/java/com/google/adk/agents/PlannerAgentTest.java +++ b/contrib/planners/src/test/java/com/google/adk/agents/PlannerAgentTest.java @@ -17,6 +17,7 @@ package com.google.adk.agents; import static com.google.common.truth.Truth.assertThat; +import static org.junit.jupiter.api.Assertions.assertThrows; import com.google.adk.events.Event; import com.google.adk.events.EventActions; @@ -25,14 +26,12 @@ import com.google.common.collect.ImmutableList; import io.reactivex.rxjava3.core.Flowable; import io.reactivex.rxjava3.core.Single; +import java.time.Instant; import java.util.List; import java.util.concurrent.atomic.AtomicInteger; -import org.junit.Test; -import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; +import org.junit.jupiter.api.Test; /** Unit tests for {@link PlannerAgent}. */ -@RunWith(JUnit4.class) public final class PlannerAgentTest { @Test @@ -86,6 +85,36 @@ public Single nextAction(PlanningContext context) { assertThat(events.get(0).content().get().text()).isEqualTo("final answer"); } + @Test + public void runAsync_withDoneWithResult_resultEventUsesInvocationProviders() { + TestBaseAgent subAgent = TestUtils.createSubAgent("sub"); + Planner resultPlanner = + new Planner() { + @Override + public Single firstAction(PlanningContext context) { + return Single.just(new PlannerAction.DoneWithResult("final answer")); + } + + @Override + public Single nextAction(PlanningContext context) { + return Single.just(new PlannerAction.Done()); + } + }; + PlannerAgent agent = + PlannerAgent.builder().name("planner").subAgents(subAgent).planner(resultPlanner).build(); + InvocationContext ctx = + TestUtils.createInvocationContext(agent).toBuilder() + .instantSource(() -> Instant.ofEpochMilli(1234L)) + .uuidProvider(() -> "fixed-uuid") + .build(); + + List events = agent.runAsync(ctx).toList().blockingGet(); + + assertThat(events).hasSize(1); + assertThat(events.get(0).id()).isEqualTo("fixed-uuid"); + assertThat(events.get(0).timestamp()).isEqualTo(1234L); + } + @Test public void runAsync_withNoOp_skipsAndContinues() { Event event1 = TestUtils.createEvent("e1"); @@ -262,10 +291,12 @@ public Single nextAction(PlanningContext context) { assertThat(events).isEmpty(); } - @Test(expected = IllegalStateException.class) + @Test public void builder_withoutPlanner_throwsIllegalState() { TestBaseAgent subAgent = TestUtils.createSubAgent("sub"); - PlannerAgent.builder().name("planner").subAgents(subAgent).build(); + PlannerAgent.Builder builder = PlannerAgent.builder().name("planner").subAgents(subAgent); + + assertThrows(IllegalStateException.class, builder::build); } @Test diff --git a/core/src/main/java/com/google/adk/agents/BaseAgent.java b/core/src/main/java/com/google/adk/agents/BaseAgent.java index c4bb0fbb7..17547c796 100644 --- a/core/src/main/java/com/google/adk/agents/BaseAgent.java +++ b/core/src/main/java/com/google/adk/agents/BaseAgent.java @@ -429,7 +429,8 @@ private Maybe callCallback( content -> { invocationContext.setEndInvocation(true); return Event.builder() - .id(Event.generateEventId()) + .id(invocationContext.newUuid()) + .timestamp(invocationContext.now().toEpochMilli()) .invocationId(invocationContext.invocationId()) .author(name()) .branch(invocationContext.branch().orElse(null)) @@ -446,7 +447,8 @@ private Maybe callCallback( if (callbackContext.state().hasDelta()) { Event.Builder eventBuilder = Event.builder() - .id(Event.generateEventId()) + .id(invocationContext.newUuid()) + .timestamp(invocationContext.now().toEpochMilli()) .invocationId(invocationContext.invocationId()) .author(name()) .branch(invocationContext.branch().orElse(null)) @@ -499,7 +501,8 @@ final Flowable checkpointAndRecord( /** Builds a resumability checkpoint event authored by this agent carrying {@code actions}. */ private Event checkpointEvent(InvocationContext context, EventActions actions) { return Event.builder() - .id(Event.generateEventId()) + .id(context.newUuid()) + .timestamp(context.now().toEpochMilli()) .invocationId(context.invocationId()) .author(name()) .branch(context.branch().orElse(null)) 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 6651c5720..1a6f1786c 100644 --- a/core/src/main/java/com/google/adk/agents/InvocationContext.java +++ b/core/src/main/java/com/google/adk/agents/InvocationContext.java @@ -27,6 +27,7 @@ import com.google.adk.flows.llmflows.Functions; import com.google.adk.memory.BaseMemoryService; import com.google.adk.models.LlmCallsLimitExceededException; +import com.google.adk.platform.UuidProvider; import com.google.adk.plugins.Plugin; import com.google.adk.plugins.PluginManager; import com.google.adk.sessions.BaseSessionService; @@ -41,6 +42,8 @@ import com.google.genai.types.FunctionResponse; import io.reactivex.rxjava3.core.Scheduler; import io.reactivex.rxjava3.schedulers.Schedulers; +import java.time.Instant; +import java.time.InstantSource; import java.util.ArrayList; import java.util.Collections; import java.util.HashSet; @@ -50,7 +53,6 @@ import java.util.Objects; import java.util.Optional; import java.util.Set; -import java.util.UUID; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.atomic.AtomicInteger; import org.jspecify.annotations.Nullable; @@ -77,6 +79,8 @@ public class InvocationContext { private final Map> agentStates; private final Map endOfAgents; @Nullable private final Scheduler scheduler; + private final InstantSource instantSource; + private final UuidProvider uuidProvider; @Nullable private String branch; private BaseAgent agent; @@ -90,7 +94,8 @@ protected InvocationContext(Builder builder) { this.liveRequestQueue = builder.liveRequestQueue; this.activeStreamingTools = builder.activeStreamingTools; this.branch = builder.branch; - this.invocationId = builder.invocationId; + // Resolved here so subclasses that construct from a builder without build() also get an id. + this.invocationId = builder.resolveInvocationId(); this.agent = builder.agent; this.session = builder.session; this.userContent = builder.userContent; @@ -107,6 +112,8 @@ protected InvocationContext(Builder builder) { this.agentStates = builder.agentStates; this.endOfAgents = builder.endOfAgents; this.scheduler = builder.scheduler; + this.instantSource = builder.instantSource; + this.uuidProvider = builder.uuidProvider; } /** Returns a new {@link Builder} for creating {@link InvocationContext} instances. */ @@ -303,9 +310,34 @@ public Scheduler scheduler() { return scheduler != null ? scheduler : Schedulers.io(); } + /** Returns the {@link InstantSource} for this invocation. */ + public InstantSource instantSource() { + return instantSource; + } + + /** Returns the {@link UuidProvider} for this invocation. */ + public UuidProvider uuidProvider() { + return uuidProvider; + } + + /** Returns the current time from this invocation's {@link InstantSource}. */ + public Instant now() { + return instantSource.instant(); + } + + /** 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(); } /** @@ -619,6 +651,7 @@ private Builder(InvocationContext context) { this.activeStreamingTools = new ConcurrentHashMap<>(context.activeStreamingTools); this.branch = context.branch; this.invocationId = context.invocationId; + this.invocationIdSet = true; this.agent = context.agent; this.session = context.session; this.userContent = context.userContent; @@ -636,6 +669,8 @@ private Builder(InvocationContext context) { this.agentStates = context.agentStates; this.endOfAgents = context.endOfAgents; this.scheduler = context.scheduler; + this.instantSource = context.instantSource; + this.uuidProvider = context.uuidProvider; } private BaseSessionService sessionService; @@ -645,7 +680,9 @@ private Builder(InvocationContext context) { @Nullable private LiveRequestQueue liveRequestQueue = null; private Map activeStreamingTools = new ConcurrentHashMap<>(); @Nullable private String branch = null; - private String invocationId = newInvocationContextId(); + private @Nullable String invocationId; + // Distinguishes "never set" (an id is minted from the UuidProvider) from an explicit null. + private boolean invocationIdSet; private BaseAgent agent; private Session session; @Nullable private Content userContent = null; @@ -659,6 +696,8 @@ private Builder(InvocationContext context) { private Map> agentStates = new ConcurrentHashMap<>(); private Map endOfAgents = new ConcurrentHashMap<>(); @Nullable private Scheduler scheduler = null; + private InstantSource instantSource = InstantSource.system(); + private UuidProvider uuidProvider = UuidProvider.SYSTEM; /** * Sets the session service for managing session state. @@ -741,6 +780,7 @@ public Builder branch(@Nullable String branch) { @CanIgnoreReturnValue public Builder invocationId(String invocationId) { this.invocationId = invocationId; + this.invocationIdSet = true; return this; } @@ -865,6 +905,44 @@ public Builder callbackContextData(Map callbackContextData) { return this; } + /** + * Sets the instant source for the invocation; {@code null} (the default) selects {@link + * InstantSource#system()}. + * + * @param instantSource the source of the current time, or {@code null} for the default. + * @return this builder instance for chaining. + */ + @CanIgnoreReturnValue + public Builder instantSource(@Nullable InstantSource instantSource) { + this.instantSource = instantSource != null ? instantSource : InstantSource.system(); + return this; + } + + /** + * Sets the UUID provider for the invocation; {@code null} (the default) selects {@link + * UuidProvider#SYSTEM}. + * + * @param uuidProvider the provider for new unique identifiers, or {@code null} for the default. + * @return this builder instance for chaining. + */ + @CanIgnoreReturnValue + public Builder uuidProvider(@Nullable UuidProvider uuidProvider) { + this.uuidProvider = uuidProvider != null ? uuidProvider : UuidProvider.SYSTEM; + return this; + } + + /** + * Returns the invocation id, minting one from the {@link UuidProvider} on first use so that + * repeated builds from this builder share it. + */ + private String resolveInvocationId() { + if (!invocationIdSet) { + invocationId = newInvocationContextId(uuidProvider); + invocationIdSet = true; + } + return invocationId; + } + /** * Builds the {@link InvocationContext} instance. * @@ -883,7 +961,8 @@ public InvocationContext build() { * @throws IllegalStateException if any required parameters are missing. */ private static void validate(Builder builder) { - if (isNullOrEmpty(builder.invocationId)) { + // An unset id is minted from the UuidProvider; an explicitly empty one is a caller error. + if (builder.invocationIdSet && isNullOrEmpty(builder.invocationId)) { throw new IllegalStateException("Invocation ID must be non-empty."); } if (builder.agent == null) { @@ -923,7 +1002,9 @@ public boolean equals(Object o) { && Objects.equals(resumabilityConfig, that.resumabilityConfig) && Objects.equals(invocationCostManager, that.invocationCostManager) && Objects.equals(callbackContextData, that.callbackContextData) - && Objects.equals(scheduler, that.scheduler); + && Objects.equals(scheduler, that.scheduler) + && Objects.equals(instantSource, that.instantSource) + && Objects.equals(uuidProvider, that.uuidProvider); } @Override @@ -947,6 +1028,8 @@ public int hashCode() { resumabilityConfig, invocationCostManager, callbackContextData, - scheduler); + scheduler, + instantSource, + uuidProvider); } } 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 40556a53f..0a654c4c1 100644 --- a/core/src/main/java/com/google/adk/events/Event.java +++ b/core/src/main/java/com/google/adk/events/Event.java @@ -23,6 +23,7 @@ import com.fasterxml.jackson.annotation.JsonProperty; import com.fasterxml.jackson.databind.annotation.JsonDeserialize; import com.google.adk.JsonBaseModel; +import com.google.adk.platform.UuidProvider; import com.google.common.collect.ImmutableList; import com.google.common.collect.Iterables; import com.google.errorprone.annotations.CanIgnoreReturnValue; @@ -34,12 +35,11 @@ import com.google.genai.types.GenerateContentResponseUsageMetadata; import com.google.genai.types.GroundingMetadata; import com.google.genai.types.Transcription; -import java.time.Instant; +import java.time.InstantSource; import java.util.List; import java.util.Objects; import java.util.Optional; import java.util.Set; -import java.util.UUID; import org.jspecify.annotations.Nullable; /** Represents an event in a session. */ @@ -71,8 +71,13 @@ public class Event extends JsonBaseModel { private Event() {} + /** + * Returns a random event id. Agents, tools, and plugins that build events inside an invocation + * should use {@code InvocationContext.newUuid()} instead so the id follows its {@link + * UuidProvider}. + */ public static String generateEventId() { - return UUID.randomUUID().toString(); + return UuidProvider.SYSTEM.newUuid(); } /** The event id. */ @@ -519,6 +524,11 @@ public Builder interrupted(@Nullable Boolean value) { return this; } + /** + * Sets the timestamp in epoch milliseconds. When unset, {@link #build()} reads the wall clock; + * events built inside an invocation should pass {@code InvocationContext.now().toEpochMilli()} + * so the timestamp follows its {@link java.time.InstantSource}. + */ @CanIgnoreReturnValue @JsonProperty("timestamp") public Builder timestamp(long value) { @@ -602,7 +612,8 @@ public Event build() { event.setCustomMetadata(customMetadata); event.setModelVersion(modelVersion); event.setActions(actions().orElseGet(() -> EventActions.builder().build())); - event.setTimestamp(timestamp().orElseGet(() -> Instant.now().toEpochMilli())); + event.setTimestamp( + timestamp().orElseGet(() -> InstantSource.system().instant().toEpochMilli())); event.setInputTranscription(inputTranscription); event.setOutputTranscription(outputTranscription); return event; @@ -641,10 +652,8 @@ public Builder toBuilder() { .customMetadata(this.customMetadata) .modelVersion(this.modelVersion) .inputTranscription(this.inputTranscription) - .outputTranscription(this.outputTranscription); - if (this.timestamp != 0) { - builder.timestamp(this.timestamp); - } + .outputTranscription(this.outputTranscription) + .timestamp(this.timestamp); return builder; } 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 d114b732e..b576a299e 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 @@ -469,7 +469,7 @@ private Flowable runOneStep(Context spanContext, InvocationContext contex resume.replayEvent(), llmRequestAfterPreprocess, spanContext) - .map(event -> event.toBuilder().id(Event.generateEventId()).build()) + .map(event -> event.toBuilder().id(context.newUuid()).build()) .concatMap(event -> followTransfer(event, context, spanContext)); } @@ -482,30 +482,23 @@ private Flowable runOneStep(Context spanContext, InvocationContext contex final Event mutableEventTemplate = Event.builder() - .id(Event.generateEventId()) + .id(context.newUuid()) + .timestamp(context.now().toEpochMilli()) .invocationId(context.invocationId()) .author(context.agent().name()) .branch(context.branch().orElse(null)) .build(); - mutableEventTemplate.setTimestamp(0L); return callLlm( spanContext, context, llmRequestAfterPreprocess, mutableEventTemplate) - .doFinally( - () -> { - String oldId = mutableEventTemplate.id(); - String newId = Event.generateEventId(); - logger.debug("Resetting event ID from {} to {}", oldId, newId); - mutableEventTemplate.setId(newId); - }) .concatMap( 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(); return followTransfer(event, context, spanContext); @@ -585,7 +578,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() @@ -663,6 +656,7 @@ public void onError(Throwable e) { Event.Builder liveEventBuilderTemplate = Event.builder() + .timestamp(invocationContext.now().toEpochMilli()) .invocationId(invocationContext.invocationId()) .author(invocationContext.agent().name()) .branch(invocationContext.branch().orElse(null)); @@ -677,7 +671,7 @@ public void onError(Throwable e) { .flatMap( llmResponse -> { Event baseEventForThisLlmResponse = - liveEventBuilderTemplate.id(Event.generateEventId()).build(); + liveEventBuilderTemplate.id(invocationContext.newUuid()).build(); return postprocess( invocationContext, baseEventForThisLlmResponse, @@ -762,7 +756,7 @@ private Flowable buildPostprocessingEvents( } Event modelResponseEvent = - buildModelResponseEvent(baseEventForLlmResponse, llmRequest, updatedResponse); + buildModelResponseEvent(context, baseEventForLlmResponse, llmRequest, updatedResponse); if (modelResponseEvent.functionCalls().isEmpty() || modelResponseEvent.partial().orElse(false)) { return processorEvents.concatWith(Flowable.just(modelResponseEvent)); @@ -872,9 +866,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)) @@ -894,7 +892,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/CodeExecution.java b/core/src/main/java/com/google/adk/flows/llmflows/CodeExecution.java index e681e48e9..91f3bddce 100644 --- a/core/src/main/java/com/google/adk/flows/llmflows/CodeExecution.java +++ b/core/src/main/java/com/google/adk/flows/llmflows/CodeExecution.java @@ -227,6 +227,8 @@ private static Flowable runPreProcessor( llmRequest.contents().add(codeContent); Event codeEvent = Event.builder() + .id(invocationContext.newUuid()) + .timestamp(invocationContext.now().toEpochMilli()) .invocationId(invocationContext.invocationId()) .author(llmAgent.name()) .content(codeContent) @@ -308,6 +310,8 @@ private static Flowable runPostProcessor( Event codeEvent = Event.builder() + .id(invocationContext.newUuid()) + .timestamp(invocationContext.now().toEpochMilli()) .invocationId(invocationContext.invocationId()) .author(llmAgent.name()) .content(responseContent) @@ -457,6 +461,8 @@ private static Single postProcessCodeExecutionResult( } eventActionsBuilder.artifactDelta(artifactDelta); return Event.builder() + .id(invocationContext.newUuid()) + .timestamp(invocationContext.now().toEpochMilli()) .invocationId(invocationContext.invocationId()) .author(invocationContext.agent().name()) .content(resultContent) 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 f11028104..a6de8132f 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 @@ -31,6 +31,7 @@ import com.google.adk.events.EventActions; import com.google.adk.events.ToolConfirmation; import com.google.adk.models.FunctionCallIds; +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; @@ -89,6 +90,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; @@ -106,7 +119,9 @@ 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(FunctionCallIds.generateClientFunctionCallId(uuidProvider)) + .build(); newParts.add(part.toBuilder().functionCall(updatedFunctionCall).build()); modified = true; } else { @@ -169,7 +184,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 +248,8 @@ public static Maybe handleFunctionCallsLive( if (events.isEmpty()) { return Maybe.empty(); } - return Maybe.fromOptional(Functions.mergeParallelFunctionResponseEvents(events)); + return Maybe.fromOptional( + Functions.mergeParallelFunctionResponseEvents(invocationContext, events)); }); } @@ -589,7 +605,7 @@ private static Maybe processFunctionResult( } private static Optional mergeParallelFunctionResponseEvents( - List functionResponseEvents) { + InvocationContext invocationContext, List functionResponseEvents) { if (functionResponseEvents.isEmpty()) { return Optional.empty(); } @@ -613,7 +629,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)) @@ -764,7 +780,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)) @@ -809,7 +826,7 @@ public static Optional generateRequestConfirmationEvent( functionCallsById.get(entry.getKey()), "toolConfirmation", entry.getValue())) - .id(generateClientFunctionCallId()) + .id(FunctionCallIds.generateClientFunctionCallId(invocationContext.uuidProvider())) .build(); longRunningToolIds.add(requestConfirmationFunctionCall.id().get()); @@ -825,6 +842,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 d1d214551..6b6ace5f2 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 @@ -115,7 +115,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/models/FunctionCallIds.java b/core/src/main/java/com/google/adk/models/FunctionCallIds.java index ba6a89b06..9451af5f8 100644 --- a/core/src/main/java/com/google/adk/models/FunctionCallIds.java +++ b/core/src/main/java/com/google/adk/models/FunctionCallIds.java @@ -16,7 +16,7 @@ package com.google.adk.models; -import java.util.UUID; +import com.google.adk.platform.UuidProvider; /** Constants and helpers for ADK-generated function call IDs. */ public final class FunctionCallIds { @@ -26,7 +26,15 @@ public final class FunctionCallIds { /** Returns a new client-side function call ID with the ADK prefix. */ public static String generateClientFunctionCallId() { - return AF_FUNCTION_CALL_ID_PREFIX + UUID.randomUUID(); + return generateClientFunctionCallId(UuidProvider.SYSTEM); + } + + /** + * Returns a new client-side function call ID with the ADK prefix, minted from the given {@link + * UuidProvider}. + */ + public static String generateClientFunctionCallId(UuidProvider uuidProvider) { + return AF_FUNCTION_CALL_ID_PREFIX + uuidProvider.newUuid(); } /** Returns whether {@code id} was generated client-side by the ADK. */ 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..ec0fe178e --- /dev/null +++ b/core/src/main/java/com/google/adk/platform/UuidProvider.java @@ -0,0 +1,34 @@ +/* + * 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. Implementations must be thread-safe and + * return a distinct value on every call; configure one with {@code + * Runner.Builder.uuidProvider(UuidProvider)}. + */ +@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..8ac03b75c --- /dev/null +++ b/core/src/main/java/com/google/adk/platform/package-info.java @@ -0,0 +1,26 @@ +/* + * 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 the system operations behind ADK-generated timestamps and ids. + * + *

An invocation reads the time from a {@link java.time.InstantSource} and mints identifiers from + * a {@link com.google.adk.platform.UuidProvider}, both configured once on the runner and carried on + * every {@code InvocationContext} it creates rather than held in static state, so concurrent + * invocations may use independent providers. Either may be called from several RxJava worker + * threads at once. + */ +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 c4af1bf75..a7b1549b1 100644 --- a/core/src/main/java/com/google/adk/runner/Runner.java +++ b/core/src/main/java/com/google/adk/runner/Runner.java @@ -41,6 +41,7 @@ import com.google.adk.flows.llmflows.PersistBarrier; import com.google.adk.memory.BaseMemoryService; import com.google.adk.models.Model; +import com.google.adk.platform.UuidProvider; import com.google.adk.plugins.Plugin; import com.google.adk.plugins.PluginManager; import com.google.adk.sessions.BaseSessionService; @@ -73,6 +74,7 @@ import io.reactivex.rxjava3.core.Scheduler; import io.reactivex.rxjava3.core.Single; import io.reactivex.rxjava3.subjects.CompletableSubject; +import java.time.InstantSource; import java.util.ArrayList; import java.util.Arrays; import java.util.HashSet; @@ -97,6 +99,8 @@ public class Runner { @Nullable private final ContextCacheConfig contextCacheConfig; private final @Nullable ResumabilityConfig resumabilityConfig; @Nullable private final Scheduler scheduler; + private final InstantSource instantSource; + private final UuidProvider uuidProvider; private final ConcurrentMap activeSessionCompletables = new MapMaker().weakValues().makeMap(); @@ -106,10 +110,14 @@ public static class Builder { private BaseAgent agent; private String appName; private BaseArtifactService artifactService = new InMemoryArtifactService(); - private BaseSessionService sessionService = new InMemorySessionService(); + private @Nullable BaseSessionService sessionService = null; + // Distinguishes "never set" (a shared in-memory default is created) from an explicit null. + private boolean sessionServiceSet; @Nullable private BaseMemoryService memoryService = null; private List plugins = ImmutableList.of(); @Nullable private Scheduler scheduler = null; + private InstantSource instantSource = InstantSource.system(); + private UuidProvider uuidProvider = UuidProvider.SYSTEM; @CanIgnoreReturnValue public Builder app(App app) { @@ -141,6 +149,7 @@ public Builder artifactService(BaseArtifactService artifactService) { @CanIgnoreReturnValue public Builder sessionService(BaseSessionService sessionService) { this.sessionService = sessionService; + this.sessionServiceSet = true; return this; } @@ -183,6 +192,32 @@ public Builder scheduler(@Nullable Scheduler scheduler) { return this; } + /** + * Sets the {@link InstantSource} that supplies every timestamp this runner generates: event + * timestamps, and the {@code lastUpdateTime} of sessions created by the default in-memory + * session service, which is created on the first {@code build()}. {@code null} (the default) + * selects {@link InstantSource#system()}. + */ + @CanIgnoreReturnValue + public Builder instantSource(@Nullable InstantSource instantSource) { + this.instantSource = instantSource != null ? instantSource : InstantSource.system(); + return this; + } + + /** + * Sets the {@link UuidProvider} that supplies every identifier this runner generates: event, + * invocation, and client function-call ids, and the ids of sessions created by the default + * in-memory session service, which is created on the first {@code build()}. {@code null} (the + * default) selects {@link UuidProvider#SYSTEM}. + */ + @CanIgnoreReturnValue + public Builder uuidProvider(@Nullable UuidProvider uuidProvider) { + this.uuidProvider = uuidProvider != null ? uuidProvider : UuidProvider.SYSTEM; + return this; + } + + // OVERWRITE is the documented default of new InMemorySessionService() until callers migrate. + @SuppressWarnings("deprecation") public Runner build() { BaseAgent buildAgent; String buildAppName; @@ -223,7 +258,15 @@ public Runner build() { throw new IllegalStateException("Artifact service must be provided."); } if (sessionService == null) { - throw new IllegalStateException("Session service must be provided."); + if (sessionServiceSet) { + throw new IllegalStateException("Session service must be provided."); + } + // Created once per builder, so every runner built from it shares the same sessions. + sessionService = + new InMemorySessionService( + InMemorySessionService.DuplicateSessionIdBehavior.OVERWRITE, + instantSource, + uuidProvider); } return new Runner( buildAgent, @@ -235,7 +278,9 @@ public Runner build() { buildEventsCompactionConfig, buildContextCacheConfig, buildResumabilityConfig, - scheduler); + scheduler, + instantSource, + uuidProvider); } } @@ -327,7 +372,9 @@ protected Runner( eventsCompactionConfig, contextCacheConfig, resumabilityConfig, - /* scheduler= */ null); + /* scheduler= */ null, + InstantSource.system(), + UuidProvider.SYSTEM); } private Runner( @@ -340,17 +387,22 @@ private Runner( @Nullable EventsCompactionConfig eventsCompactionConfig, @Nullable ContextCacheConfig contextCacheConfig, @Nullable ResumabilityConfig resumabilityConfig, - @Nullable Scheduler scheduler) { + @Nullable Scheduler scheduler, + InstantSource instantSource, + UuidProvider uuidProvider) { this.agent = agent; this.appName = appName; this.artifactService = artifactService; this.sessionService = sessionService; this.memoryService = memoryService; this.pluginManager = new PluginManager(plugins); - this.eventsCompactionConfig = createEventsCompactionConfig(agent, eventsCompactionConfig); + this.eventsCompactionConfig = + createEventsCompactionConfig(agent, eventsCompactionConfig, instantSource, uuidProvider); this.contextCacheConfig = contextCacheConfig; this.resumabilityConfig = resumabilityConfig; this.scheduler = scheduler; + this.instantSource = instantSource; + this.uuidProvider = uuidProvider; } /** @@ -462,7 +514,8 @@ private Single appendNewMessageToSession( } Event.Builder eventBuilder = Event.builder() - .id(Event.generateEventId()) + .id(invocationContext.newUuid()) + .timestamp(invocationContext.now().toEpochMilli()) .invocationId(invocationContext.invocationId()) .author(Role.USER) .branch(branch) @@ -685,7 +738,7 @@ private Flowable runNewInvocation( () -> { Context capturedContext = Context.current(); BaseAgent rootAgent = this.agent; - String invocationId = InvocationContext.newInvocationContextId(); + String invocationId = InvocationContext.newInvocationContextId(this.uuidProvider); // Eagerly merge stateDelta into session state so onUserMessageCallback() can see it. // Safe: this in-memory update is ahead of time; canonical persistence still happens @@ -859,7 +912,8 @@ private Maybe beforeRunEventFor(InvocationContext context) { .map( content -> Event.builder() - .id(Event.generateEventId()) + .id(context.newUuid()) + .timestamp(context.now().toEpochMilli()) .invocationId(context.invocationId()) .author("model") .content(content) @@ -981,7 +1035,8 @@ private Flowable resumeCore( // No message to carry the delta, so persist a content-less event, as Python does. Event stateDeltaEvent = Event.builder() - .id(Event.generateEventId()) + .id(initialContext.newUuid()) + .timestamp(initialContext.now().toEpochMilli()) .invocationId(resolvedInvocationId) .author(Role.USER) .actions( @@ -1323,6 +1378,8 @@ private InvocationContext.Builder newInvocationContextBuilder(Session session, B .contextCacheConfig(this.contextCacheConfig) .resumabilityConfig(this.resumabilityConfig) .scheduler(this.scheduler) + .instantSource(this.instantSource) + .uuidProvider(this.uuidProvider) .agent(agent); } @@ -1538,7 +1595,10 @@ private boolean hasLiveRequestQueueParameter(FunctionTool functionTool) { @Nullable private static EventsCompactionConfig createEventsCompactionConfig( - BaseAgent agent, @Nullable EventsCompactionConfig config) { + BaseAgent agent, + @Nullable EventsCompactionConfig config, + InstantSource instantSource, + UuidProvider uuidProvider) { if (config == null || config.summarizer() != null) { return config; } @@ -1548,7 +1608,10 @@ private static EventsCompactionConfig createEventsCompactionConfig( .map(LlmAgent.class::cast) .flatMap(LlmAgent::model) .flatMap(Model::model) - .map(LlmEventSummarizer::new) + .map( + baseLlm -> + new LlmEventSummarizer( + baseLlm, /* promptTemplate= */ null, instantSource, uuidProvider)) .orElseThrow( () -> new IllegalArgumentException( 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 a9874f238..fa7c19193 100644 --- a/core/src/main/java/com/google/adk/sessions/InMemorySessionService.java +++ b/core/src/main/java/com/google/adk/sessions/InMemorySessionService.java @@ -20,19 +20,20 @@ import com.google.adk.events.Event; import com.google.adk.events.EventActions; +import com.google.adk.platform.UuidProvider; import com.google.common.collect.ImmutableList; import com.google.errorprone.annotations.CanIgnoreReturnValue; import io.reactivex.rxjava3.core.Completable; import io.reactivex.rxjava3.core.Maybe; import io.reactivex.rxjava3.core.Single; import java.time.Instant; +import java.time.InstantSource; import java.util.ArrayList; import java.util.Collection; import java.util.List; 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; @@ -79,6 +80,8 @@ public enum DuplicateSessionIdBehavior { private final ConcurrentMap> appState; private final DuplicateSessionIdBehavior duplicateSessionIdBehavior; + private final InstantSource instantSource; + private final UuidProvider uuidProvider; /** * Creates a new instance of the in-memory session service with empty storage. A duplicate session @@ -95,12 +98,29 @@ public InMemorySessionService() { * already in use. */ public InMemorySessionService(DuplicateSessionIdBehavior duplicateSessionIdBehavior) { + this(duplicateSessionIdBehavior, InstantSource.system(), UuidProvider.SYSTEM); + } + + /** + * Creates a new instance of the in-memory session service with empty storage. + * + * @param duplicateSessionIdBehavior What {@code createSession} does when the session ID is + * already in use. + * @param instantSource Supplies the {@code lastUpdateTime} of a newly created session. + * @param uuidProvider Supplies the session ID when {@code createSession} is not given one. + */ + public InMemorySessionService( + DuplicateSessionIdBehavior duplicateSessionIdBehavior, + InstantSource instantSource, + UuidProvider uuidProvider) { this.duplicateSessionIdBehavior = Objects.requireNonNull( duplicateSessionIdBehavior, "duplicateSessionIdBehavior cannot be null"); this.sessions = new ConcurrentHashMap<>(); this.userState = new ConcurrentHashMap<>(); this.appState = new ConcurrentHashMap<>(); + this.instantSource = Objects.requireNonNull(instantSource, "instantSource cannot be null"); + this.uuidProvider = Objects.requireNonNull(uuidProvider, "uuidProvider cannot be null"); } @Override @@ -125,7 +145,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 = @@ -137,7 +157,7 @@ public Single createSession( .appName(appName) .userId(userId) .state(initialState) - .lastUpdateTime(Instant.now()) + .lastUpdateTime(instantSource.instant()) .build(); ConcurrentMap userSessions = diff --git a/core/src/main/java/com/google/adk/summarizer/LlmEventSummarizer.java b/core/src/main/java/com/google/adk/summarizer/LlmEventSummarizer.java index 7d1c8a821..a57c701eb 100644 --- a/core/src/main/java/com/google/adk/summarizer/LlmEventSummarizer.java +++ b/core/src/main/java/com/google/adk/summarizer/LlmEventSummarizer.java @@ -26,6 +26,7 @@ import com.google.adk.events.EventCompaction; import com.google.adk.models.BaseLlm; import com.google.adk.models.LlmRequest; +import com.google.adk.platform.UuidProvider; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.genai.types.Content; @@ -33,8 +34,11 @@ import com.google.genai.types.FunctionResponse; import com.google.genai.types.Part; import io.reactivex.rxjava3.core.Maybe; +import java.time.InstantSource; import java.util.List; +import java.util.Objects; import java.util.Optional; +import org.jspecify.annotations.Nullable; /** An LLM-based event summarizer for sliding window compaction. */ public final class LlmEventSummarizer implements BaseEventSummarizer { @@ -52,14 +56,33 @@ public final class LlmEventSummarizer implements BaseEventSummarizer { private final BaseLlm baseLlm; private final String promptTemplate; + private final InstantSource instantSource; + private final UuidProvider uuidProvider; public LlmEventSummarizer(BaseLlm baseLlm) { this(baseLlm, DEFAULT_PROMPT_TEMPLATE); } - public LlmEventSummarizer(BaseLlm baseLlm, String promptTemplate) { + /** + * Creates a summarizer with the given prompt template; {@code null} selects the default template. + */ + public LlmEventSummarizer(BaseLlm baseLlm, @Nullable String promptTemplate) { + this(baseLlm, promptTemplate, InstantSource.system(), UuidProvider.SYSTEM); + } + + /** + * Creates a summarizer whose compaction events draw their id, invocation id, and timestamp from + * the given sources. Passing {@code null} as {@code promptTemplate} selects the default template. + */ + public LlmEventSummarizer( + BaseLlm baseLlm, + @Nullable String promptTemplate, + InstantSource instantSource, + UuidProvider uuidProvider) { this.baseLlm = baseLlm; - this.promptTemplate = promptTemplate; + this.promptTemplate = promptTemplate != null ? promptTemplate : DEFAULT_PROMPT_TEMPLATE; + this.instantSource = Objects.requireNonNull(instantSource, "instantSource cannot be null"); + this.uuidProvider = Objects.requireNonNull(uuidProvider, "uuidProvider cannot be null"); } @Override @@ -101,10 +124,11 @@ public Maybe summarizeEvents(List events) { .map( compaction -> Event.builder() - .id(Event.generateEventId()) + .id(uuidProvider.newUuid()) + .timestamp(instantSource.instant().toEpochMilli()) .author(Role.USER) .actions(EventActions.builder().compaction(compaction).build()) - .invocationId(Event.generateEventId()) + .invocationId(uuidProvider.newUuid()) .build()))); } diff --git a/core/src/main/java/com/google/adk/tools/AgentTool.java b/core/src/main/java/com/google/adk/tools/AgentTool.java index 86a5eb44f..bb55922b4 100644 --- a/core/src/main/java/com/google/adk/tools/AgentTool.java +++ b/core/src/main/java/com/google/adk/tools/AgentTool.java @@ -24,6 +24,7 @@ import com.google.adk.agents.BaseAgentConfig; import com.google.adk.agents.ConfigAgentUtils; import com.google.adk.agents.ConfigAgentUtils.ConfigurationException; +import com.google.adk.agents.InvocationContext; import com.google.adk.agents.LlmAgent; import com.google.adk.agents.RunConfig; import com.google.adk.agents.RunConfig.StreamingMode; @@ -183,21 +184,22 @@ public Single> runAsync(Map args, ToolContex content = Content.fromParts(Part.fromText(input.toString())); } + InvocationContext parentContext = toolContext.invocationContext(); ImmutableList plugins = - this.includePlugins - ? ImmutableList.of(toolContext.invocationContext().pluginManager()) - : ImmutableList.of(); - // The nested run inherits the caller's scheduler so its fan-out stays on the same scheduler. + this.includePlugins ? ImmutableList.of(parentContext.pluginManager()) : ImmutableList.of(); + // The nested run inherits the caller's scheduler, clock, and id provider. Runner runner = Runner.builder() .agent(this.agent) .appName(toolContext.agentName()) .plugins(plugins) .memoryService(new InMemoryMemoryService()) - .scheduler(toolContext.invocationContext().scheduler()) + .scheduler(parentContext.scheduler()) + .instantSource(parentContext.instantSource()) + .uuidProvider(parentContext.uuidProvider()) .build(); // Follow the caller's RunConfig but run unary: only the last event becomes the result. - RunConfig callerRunConfig = toolContext.invocationContext().runConfig(); + RunConfig callerRunConfig = parentContext.runConfig(); RunConfig runConfig = callerRunConfig.streamingMode() == StreamingMode.NONE ? callerRunConfig 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 adcf49ab7..8e19ce355 100644 --- a/core/src/test/java/com/google/adk/agents/InvocationContextTest.java +++ b/core/src/test/java/com/google/adk/agents/InvocationContextTest.java @@ -26,6 +26,7 @@ import com.google.adk.events.EventActions; import com.google.adk.memory.BaseMemoryService; import com.google.adk.models.LlmCallsLimitExceededException; +import com.google.adk.platform.UuidProvider; import com.google.adk.plugins.PluginManager; import com.google.adk.sessions.BaseSessionService; import com.google.adk.sessions.Session; @@ -39,9 +40,12 @@ import com.google.genai.types.Part; import io.reactivex.rxjava3.core.Scheduler; import io.reactivex.rxjava3.schedulers.Schedulers; +import java.time.Instant; +import java.time.InstantSource; import java.util.HashMap; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicInteger; import org.jspecify.annotations.Nullable; import org.junit.Assert; import org.junit.Before; @@ -1489,4 +1493,182 @@ public void populateInvocationAgentStates_noContentNoState_ignored() { assertThat(context.agentStates()).isEmpty(); assertThat(context.endOfAgents()).isEmpty(); } + + @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.instantSource()).isEqualTo(InstantSource.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_nullSelectsTheDefaults() { + InvocationContext context = + InvocationContext.builder() + .sessionService(mockSessionService) + .artifactService(mockArtifactService) + .agent(mockAgent) + .session(session) + .instantSource(() -> Instant.ofEpochMilli(1234L)) + .uuidProvider(() -> "fixed-uuid") + .instantSource(null) + .uuidProvider(null) + .build(); + + assertThat(context.instantSource()).isEqualTo(InstantSource.system()); + assertThat(context.uuidProvider()).isEqualTo(UuidProvider.SYSTEM); + } + + @Test + public void equals_includesTheProviders() { + InvocationContext defaults = + InvocationContext.builder() + .sessionService(mockSessionService) + .artifactService(mockArtifactService) + .agent(mockAgent) + .session(session) + .invocationId("same-id") + .build(); + InvocationContext custom = + defaults.toBuilder() + .instantSource(() -> Instant.ofEpochMilli(1234L)) + .uuidProvider(() -> "fixed-uuid") + .build(); + + assertThat(defaults.toBuilder().build()).isEqualTo(defaults); + assertThat(defaults.toBuilder().build().hashCode()).isEqualTo(defaults.hashCode()); + assertThat(custom.toBuilder().build()).isEqualTo(custom); + assertThat(custom).isNotEqualTo(defaults); + } + + @Test + public void providers_customAreCarriedAndUsedByAccessors() { + Instant fixed = Instant.ofEpochMilli(1234L); + InstantSource fixedClock = () -> fixed; + UuidProvider fixedUuids = () -> "fixed-uuid"; + + InvocationContext context = + InvocationContext.builder() + .sessionService(mockSessionService) + .artifactService(mockArtifactService) + .agent(mockAgent) + .session(session) + .instantSource(fixedClock) + .uuidProvider(fixedUuids) + .build(); + + assertThat(context.instantSource()).isEqualTo(fixedClock); + assertThat(context.uuidProvider()).isEqualTo(fixedUuids); + assertThat(context.now()).isEqualTo(fixed); + assertThat(context.newUuid()).isEqualTo("fixed-uuid"); + } + + @Test + public void toBuilder_carriesProviders() { + InstantSource fixedClock = () -> Instant.ofEpochMilli(1234L); + UuidProvider fixedUuids = () -> "fixed-uuid"; + + InvocationContext originalContext = + InvocationContext.builder() + .sessionService(mockSessionService) + .artifactService(mockArtifactService) + .agent(mockAgent) + .session(session) + .instantSource(fixedClock) + .uuidProvider(fixedUuids) + .build(); + + InvocationContext copiedContext = originalContext.toBuilder().build(); + + assertThat(copiedContext.instantSource()).isEqualTo(fixedClock); + assertThat(copiedContext.uuidProvider()).isEqualTo(fixedUuids); + } + + @Test + public void build_withoutInvocationId_mintsIdFromUuidProvider() { + InvocationContext context = + InvocationContext.builder() + .sessionService(mockSessionService) + .artifactService(mockArtifactService) + .agent(mockAgent) + .session(session) + .uuidProvider(() -> "fixed-uuid") + .build(); + + assertThat(context.invocationId()).isEqualTo("e-fixed-uuid"); + } + + @Test + public void toBuilder_keepsLazilyMintedInvocationId_withoutDrawingAgain() { + AtomicInteger draws = new AtomicInteger(); + InvocationContext original = + InvocationContext.builder() + .sessionService(mockSessionService) + .artifactService(mockArtifactService) + .agent(mockAgent) + .session(session) + .uuidProvider(() -> "uuid-" + draws.incrementAndGet()) + .build(); + + InvocationContext copy = original.toBuilder().build(); + + assertThat(original.invocationId()).isEqualTo("e-uuid-1"); + assertThat(copy.invocationId()).isEqualTo(original.invocationId()); + assertThat(draws.get()).isEqualTo(1); + } + + @Test + public void build_twiceWithoutInvocationId_sharesOneMintedId() { + AtomicInteger draws = new AtomicInteger(); + InvocationContext.Builder builder = + InvocationContext.builder() + .sessionService(mockSessionService) + .artifactService(mockArtifactService) + .agent(mockAgent) + .session(session) + .uuidProvider(() -> "uuid-" + draws.incrementAndGet()); + + InvocationContext first = builder.build(); + InvocationContext second = builder.build(); + + assertThat(first.invocationId()).isEqualTo("e-uuid-1"); + assertThat(second.invocationId()).isEqualTo("e-uuid-1"); + assertThat(draws.get()).isEqualTo(1); + } + + @Test + public void build_withExplicitInvocationId_keepsItRegardlessOfUuidProvider() { + InvocationContext context = + InvocationContext.builder() + .sessionService(mockSessionService) + .artifactService(mockArtifactService) + .agent(mockAgent) + .session(session) + .invocationId("explicit-id") + .uuidProvider(() -> "fixed-uuid") + .build(); + + assertThat(context.invocationId()).isEqualTo("explicit-id"); + } } 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 da8aa5eb5..52e09b071 100644 --- a/core/src/test/java/com/google/adk/events/EventTest.java +++ b/core/src/test/java/com/google/adk/events/EventTest.java @@ -124,6 +124,21 @@ public void event_builder_fills_default_timestamp() { assertThat(event.timestamp()).isAtMost(after); } + @Test + public void event_toBuilder_preserves_zero_timestamp() { + Event event = + Event.builder() + .id("event_id") + .invocationId("invocation_id") + .author("agent") + .timestamp(0L) + .build(); + + Event copy = event.toBuilder().build(); + + assertThat(copy.timestamp()).isEqualTo(0L); + } + @Test public void event_equals_works() { Event event1 = diff --git a/core/src/test/java/com/google/adk/flows/llmflows/CodeExecutionTest.java b/core/src/test/java/com/google/adk/flows/llmflows/CodeExecutionTest.java index 7e8882939..140f8f7f9 100644 --- a/core/src/test/java/com/google/adk/flows/llmflows/CodeExecutionTest.java +++ b/core/src/test/java/com/google/adk/flows/llmflows/CodeExecutionTest.java @@ -46,7 +46,7 @@ import com.google.genai.types.Content; import com.google.genai.types.Part; import io.reactivex.rxjava3.core.Single; -import io.reactivex.rxjava3.observers.TestObserver; +import java.time.Instant; import java.util.ArrayList; import org.junit.Before; import org.junit.Rule; @@ -81,6 +81,8 @@ public void setUp() { when(invocationContext.appName()).thenReturn("app"); when(invocationContext.userId()).thenReturn("user"); when(invocationContext.artifactService()).thenReturn(mockArtifactService); + when(invocationContext.now()).thenReturn(Instant.ofEpochMilli(1234L)); + when(invocationContext.newUuid()).thenReturn("test-event-id"); when(mockArtifactService.saveArtifact( anyString(), anyString(), anyString(), anyString(), any(Part.class))) .thenReturn(Single.just(1)); @@ -116,6 +118,11 @@ public void testResponseProcessor_withCode_executesCode() { ImmutableList events = ImmutableList.copyOf(result.events()); assertThat(events).hasSize(2); + // Both code-execution events draw their id and timestamp from the invocation context. + for (Event event : events) { + assertThat(event.id()).isEqualTo("test-event-id"); + assertThat(event.timestamp()).isEqualTo(1234L); + } Part executableCodePart = events.get(0).content().get().parts().get().get(1); assertThat(executableCodePart.executableCode().get().code()).hasValue(code); @@ -181,11 +188,19 @@ public void testRequestProcessor_withCode_hasNoErrors() throws Exception { .build()))); // act - Single result = - CodeExecution.requestProcessor.processRequest(invocationContext, llmReqBuilder.build()); - TestObserver testObserver = result.test(); + RequestProcessingResult result = + CodeExecution.requestProcessor + .processRequest(invocationContext, llmReqBuilder.build()) + .blockingGet(); // assert - testObserver.assertNoErrors(); + ImmutableList events = ImmutableList.copyOf(result.events()); + assertThat(events).hasSize(2); + // The data-file processing event and its execution result draw their id and timestamp from + // the invocation context. + for (Event event : events) { + assertThat(event.id()).isEqualTo("test-event-id"); + assertThat(event.timestamp()).isEqualTo(1234L); + } } } 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 0d3d99fbc..447726685 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 @@ -44,6 +44,8 @@ import com.google.genai.types.Part; import io.reactivex.rxjava3.core.Single; import io.reactivex.rxjava3.schedulers.Schedulers; +import java.time.Instant; +import java.time.InstantSource; import java.util.ArrayList; import java.util.LinkedHashMap; import java.util.List; @@ -563,26 +565,7 @@ private static int workersCreatedForParallelSubscribe(BaseAgent agent) { } private static Event twoToolCallsEvent() { - return createEvent("event").toBuilder() - .content( - Content.fromParts( - Part.builder() - .functionCall( - FunctionCall.builder() - .id("call_1") - .name("tool_1") - .args(ImmutableMap.of()) - .build()) - .build(), - Part.builder() - .functionCall( - FunctionCall.builder() - .id("call_2") - .name("tool_2") - .args(ImmutableMap.of()) - .build()) - .build())) - .build(); + return toolCallsEvent("tool_1", "tool_2"); } /** Single-tool case bypasses the parallel scheduler path; must still return the correct event. */ @@ -783,4 +766,83 @@ public Single> runAsync(Map args, ToolContex }); } } + + @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_singleCall_responseEventUsesInvocationProviders() { + InvocationContext invocationContext = providerContext(); + + Event responseEvent = + Functions.handleFunctionCalls( + invocationContext, + toolCallsEvent("tool_1"), + ImmutableMap.of("tool_1", new SleepingTool("tool_1", 0L))) + .blockingGet(); + + assertThat(responseEvent.id()).isEqualTo("fixed-uuid"); + assertThat(responseEvent.timestamp()).isEqualTo(0L); + assertThat(responseEvent.functionResponses()).hasSize(1); + } + + @Test + public void handleFunctionCalls_parallelCalls_mergedEventUsesInvocationProviders() { + InvocationContext invocationContext = providerContext(); + + Event mergedEvent = + Functions.handleFunctionCalls( + invocationContext, + toolCallsEvent("tool_1", "tool_2"), + ImmutableMap.of( + "tool_1", + new SleepingTool("tool_1", 0L), + "tool_2", + new SleepingTool("tool_2", 0L))) + .blockingGet(); + + assertThat(mergedEvent.id()).isEqualTo("fixed-uuid"); + assertThat(mergedEvent.timestamp()).isEqualTo(0L); + assertThat(mergedEvent.functionResponses()).hasSize(2); + } + + /** A context whose clock is fixed at the epoch and whose ids are all {@code fixed-uuid}. */ + private static InvocationContext providerContext() { + return createInvocationContext(createRootAgent()).toBuilder() + .instantSource(InstantSource.fixed(Instant.EPOCH)) + .uuidProvider(() -> "fixed-uuid") + .build(); + } + + /** A model event calling each named tool once, with model-assigned call ids. */ + private static Event toolCallsEvent(String... toolNames) { + ImmutableList.Builder parts = ImmutableList.builder(); + for (int i = 0; i < toolNames.length; i++) { + parts.add( + Part.builder() + .functionCall( + FunctionCall.builder() + .id("call_" + (i + 1)) + .name(toolNames[i]) + .args(ImmutableMap.of()) + .build()) + .build()); + } + return createEvent("event").toBuilder() + .content(Content.builder().role("model").parts(parts.build()).build()) + .build(); + } } diff --git a/core/src/test/java/com/google/adk/flows/llmflows/OutputSchemaTest.java b/core/src/test/java/com/google/adk/flows/llmflows/OutputSchemaTest.java index d9dbfed0a..20d47dca6 100644 --- a/core/src/test/java/com/google/adk/flows/llmflows/OutputSchemaTest.java +++ b/core/src/test/java/com/google/adk/flows/llmflows/OutputSchemaTest.java @@ -45,6 +45,8 @@ import com.google.genai.types.Schema; import io.reactivex.rxjava3.core.Flowable; import io.reactivex.rxjava3.core.Single; +import java.time.Instant; +import java.time.InstantSource; import java.util.List; import java.util.Map; import java.util.Objects; @@ -86,6 +88,22 @@ public Single> runAsync(Map args, ToolContex } } + @Test + public void createFinalModelResponseEvent_usesInvocationProviders() { + LlmAgent agent = LlmAgent.builder().name("agent").model(testLlm).build(); + InvocationContext context = + createInvocationContext(agent).toBuilder() + .instantSource(InstantSource.fixed(Instant.ofEpochMilli(1234L))) + .uuidProvider(() -> "fixed-uuid") + .build(); + + Event event = OutputSchema.createFinalModelResponseEvent(context, "{\"answer\": 42}"); + + assertThat(event.id()).isEqualTo("fixed-uuid"); + assertThat(event.timestamp()).isEqualTo(1234L); + assertThat(event.invocationId()).isEqualTo(context.invocationId()); + } + @Test public void processRequest_noOutputSchema_doesNothing() { LlmAgent agent = diff --git a/core/src/test/java/com/google/adk/models/FunctionCallIdsTest.java b/core/src/test/java/com/google/adk/models/FunctionCallIdsTest.java index a56878e14..c2f01e028 100644 --- a/core/src/test/java/com/google/adk/models/FunctionCallIdsTest.java +++ b/core/src/test/java/com/google/adk/models/FunctionCallIdsTest.java @@ -17,6 +17,7 @@ import static com.google.common.truth.Truth.assertThat; +import com.google.adk.platform.UuidProvider; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -39,6 +40,16 @@ public void generateClientFunctionCallId_returnsUniqueIds() { assertThat(first).isNotEqualTo(second); } + @Test + public void generateClientFunctionCallId_withProvider_derivesFromProvider() { + UuidProvider provider = () -> "deterministic-uuid"; + + String id = FunctionCallIds.generateClientFunctionCallId(provider); + + assertThat(id).isEqualTo("adk-deterministic-uuid"); + assertThat(FunctionCallIds.isClientGeneratedFunctionCallId(id)).isTrue(); + } + @Test public void isClientGeneratedFunctionCallId_falseForModelGeneratedId() { assertThat(FunctionCallIds.isClientGeneratedFunctionCallId("call_123")).isFalse(); 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/RunnerResumabilityTest.java b/core/src/test/java/com/google/adk/runner/RunnerResumabilityTest.java index c612c52e5..147d4636d 100644 --- a/core/src/test/java/com/google/adk/runner/RunnerResumabilityTest.java +++ b/core/src/test/java/com/google/adk/runner/RunnerResumabilityTest.java @@ -56,6 +56,7 @@ import com.google.adk.agents.RunConfig; import com.google.adk.agents.SequentialAgent; import com.google.adk.apps.App; +import com.google.adk.apps.ResumabilityConfig; import com.google.adk.events.Event; import com.google.adk.events.EventActions; import com.google.adk.flows.llmflows.Functions; @@ -78,6 +79,7 @@ import io.opentelemetry.sdk.testing.junit4.OpenTelemetryRule; import io.reactivex.rxjava3.core.Flowable; import io.reactivex.rxjava3.core.Maybe; +import java.time.Instant; import java.util.List; import java.util.Objects; import java.util.Optional; @@ -976,6 +978,68 @@ public void resume_withStateDeltaAndNoMessage_persistsStateDeltaEvent() { .isTrue(); } + // Checkpoint (end-of-agent) events and the content-less stateDelta event minted on resume draw + // their ids and timestamps from the Runner's providers, like every other persisted event. + @Test + public void resumableRunner_withProviders_checkpointAndStateDeltaEventsUseProviders() { + TestLlm testLlm = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "draft")), + createTextLlmResponse("done")); + LlmAgent agent = + createTestAgentBuilder(testLlm).name("agent").tools(pendingFunctionTool()).build(); + AtomicInteger counter = new AtomicInteger(); + Runner runner = + Runner.builder() + .app( + App.builder() + .name("test") + .rootAgent(agent) + .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) + .build()) + .instantSource(() -> Instant.ofEpochMilli(1234L)) + .uuidProvider(() -> String.format("uuid-%04d", counter.getAndIncrement())) + .build(); + Session session = newSession(runner); + ImmutableList turn1 = runTurn(runner, session, "start"); + String invocationId = turn1.get(0).invocationId(); + var unused = + resume(runner, session, invocationId, /* newMessage= */ null, ImmutableMap.of("key1", "v")); + // A second, completing turn records an end-of-agent checkpoint. + Runner completingRunner = + Runner.builder() + .app( + App.builder() + .name("test") + .rootAgent(textAgent("agent", "hello")) + .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) + .build()) + .sessionService(runner.sessionService()) + .instantSource(() -> Instant.ofEpochMilli(1234L)) + .uuidProvider(() -> String.format("uuid-%04d", counter.getAndIncrement())) + .build(); + ImmutableList turn2 = runTurn(completingRunner, session, "again"); + + assertThat(session.id()).isEqualTo("uuid-0000"); + assertThat(invocationId).isEqualTo("e-uuid-0001"); + assertEndOfAgent(turn2, "agent"); + Session reloaded = reloadSession(runner, session); + assertThat( + reloaded.events().stream() + .anyMatch( + event -> + event.content().isEmpty() + && event.actions().stateDelta().containsKey("key1"))) + .isTrue(); + for (Event event : reloaded.events()) { + assertThat(event.id()).startsWith("uuid-"); + assertThat(event.timestamp()).isEqualTo(1234L); + } + assertThat(reloaded.events().stream().map(Event::id).distinct().count()) + .isEqualTo(reloaded.events().size()); + } + // ResumeInvocationTest parity: a function response resumes an older invocation, not the latest. @Test public void resume_resumesAnyInvocation_notJustTheLatest() { 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 a38502cbb..c92f67ff3 100644 --- a/core/src/test/java/com/google/adk/runner/RunnerTest.java +++ b/core/src/test/java/com/google/adk/runner/RunnerTest.java @@ -35,10 +35,12 @@ import static com.google.adk.testing.TestUtils.simplifyEvents; import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.truth.Truth.assertThat; +import static com.google.common.truth.Truth.assertWithMessage; import static java.nio.charset.StandardCharsets.UTF_8; import static java.util.Arrays.stream; import static java.util.concurrent.TimeUnit.MILLISECONDS; import static java.util.concurrent.TimeUnit.SECONDS; +import static org.junit.Assert.assertThrows; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.ArgumentMatchers.eq; @@ -64,8 +66,11 @@ import com.google.adk.artifacts.InMemoryArtifactService; import com.google.adk.events.Event; import com.google.adk.flows.llmflows.Functions; +import com.google.adk.models.BaseLlm; +import com.google.adk.models.BaseLlmConnection; import com.google.adk.models.LlmRequest; import com.google.adk.models.LlmResponse; +import com.google.adk.platform.UuidProvider; import com.google.adk.plugins.BasePlugin; import com.google.adk.sessions.BaseSessionService; import com.google.adk.sessions.GetSessionConfig; @@ -83,6 +88,7 @@ import com.google.adk.testing.TestUtils.FailingEchoTool; import com.google.adk.tools.BaseTool; import com.google.adk.tools.FunctionTool; +import com.google.adk.tools.SetModelResponseTool; import com.google.adk.tools.ToolContext; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; @@ -94,6 +100,7 @@ import com.google.genai.types.FunctionResponse; import com.google.genai.types.Part; import com.google.genai.types.PartialArg; +import com.google.genai.types.Schema; import io.opentelemetry.api.trace.Tracer; import io.opentelemetry.context.Context; import io.opentelemetry.context.ContextKey; @@ -108,6 +115,7 @@ import io.reactivex.rxjava3.subjects.PublishSubject; import io.reactivex.rxjava3.subscribers.TestSubscriber; import java.time.Instant; +import java.time.InstantSource; import java.util.ArrayList; import java.util.Collections; import java.util.List; @@ -3535,4 +3543,409 @@ public Single listEvents(String appName, String userId, Stri return delegate.listEvents(appName, userId, sessionId); } } + + /** + * 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_producesIdenticalEvents() { + 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.stream().distinct().toList()).containsExactly(1234L); + // The invocation id is minted from the injected UUID provider. + assertThat(firstInvocationIds.stream().distinct().toList()).containsExactly("e-uuid-0000"); + } + + @Test + public void build_defaultSessionService_isSharedByRunnersFromTheSameBuilder() { + Runner.Builder builder = + Runner.builder().app(App.builder().name("test").rootAgent(agent).build()); + + Runner first = builder.build(); + Runner second = builder.build(); + Session created = + first.sessionService().createSession("test", "user", null, "shared-session").blockingGet(); + + assertThat(second.sessionService()).isSameInstanceAs(first.sessionService()); + assertThat( + second + .sessionService() + .getSession("test", "user", created.id(), Optional.empty()) + .blockingGet()) + .isNotNull(); + } + + @Test + public void build_explicitSessionService_isUsedAsIs() { + BaseSessionService explicitSessionService = new InMemorySessionService(); + + Runner built = + Runner.builder() + .app(App.builder().name("test").rootAgent(agent).build()) + .sessionService(explicitSessionService) + .build(); + + assertThat(built.sessionService()).isSameInstanceAs(explicitSessionService); + } + + @Test + public void build_withNullSessionService_throws() { + Runner.Builder builder = + Runner.builder() + .app(App.builder().name("test").rootAgent(agent).build()) + .sessionService(null); + + IllegalStateException exception = assertThrows(IllegalStateException.class, builder::build); + + assertThat(exception).hasMessageThat().isEqualTo("Session service must be provided."); + } + + @Test + public void build_withNullThenSessionService_usesTheSessionService() { + BaseSessionService explicitSessionService = new InMemorySessionService(); + + Runner built = + Runner.builder() + .app(App.builder().name("test").rootAgent(agent).build()) + .sessionService(null) + .sessionService(explicitSessionService) + .build(); + + assertThat(built.sessionService()).isSameInstanceAs(explicitSessionService); + } + + @Test + public void builder_nullProviders_selectTheSystemDefaults() { + Runner nullProviderRunner = + Runner.builder() + .app(App.builder().name("test").rootAgent(agent).build()) + .instantSource(InstantSource.fixed(Instant.EPOCH)) + .uuidProvider(() -> "fixed-uuid") + .instantSource(null) + .uuidProvider(null) + .build(); + + Session created = + nullProviderRunner.sessionService().createSession("test", "user").blockingGet(); + + assertThat(created.id()).isNotEqualTo("fixed-uuid"); + assertThat(created.lastUpdateTime()).isGreaterThan(Instant.EPOCH); + } + + @Test + public void build_defaultSessionService_usesRunnerProviders() { + InstantSource fixedClock = () -> Instant.ofEpochMilli(1234L); + UuidProvider fixedUuids = () -> "session-uuid"; + Runner deterministicRunner = + Runner.builder() + .app(App.builder().name("test").rootAgent(agent).build()) + .instantSource(fixedClock) + .uuidProvider(fixedUuids) + .build(); + + Session createdSession = + deterministicRunner.sessionService().createSession("test", "user").blockingGet(); + + assertThat(createdSession.id()).isEqualTo("session-uuid"); + assertThat(createdSession.lastUpdateTime()).isEqualTo(Instant.ofEpochMilli(1234L)); + } + + @Test + public void beforeRunCallback_eventUsesRunnerProviders() { + BasePlugin providerPlugin = mockPlugin("provider-plugin"); + when(providerPlugin.beforeRunCallback(any())).thenReturn(Maybe.just(pluginContent)); + Runner deterministicRunner = + Runner.builder() + .app( + App.builder() + .name("test") + .rootAgent(agent) + .plugins(ImmutableList.of(providerPlugin)) + .build()) + .instantSource(() -> Instant.ofEpochMilli(1234L)) + .uuidProvider(() -> "plugin-event-uuid") + .build(); + Session deterministicSession = + deterministicRunner.sessionService().createSession("test", "user").blockingGet(); + + List events = + deterministicRunner + .runAsync("user", deterministicSession.id(), createContent("will not be processed")) + .toList() + .blockingGet(); + + assertThat(simplifyEvents(events)).containsExactly("model: from plugin"); + assertThat(events.get(0).id()).isEqualTo("plugin-event-uuid"); + assertThat(events.get(0).timestamp()).isEqualTo(1234L); + } + + @Test + public void runLive_withProviders_eventsUseRunnerProviders() throws Exception { + AtomicInteger counter = new AtomicInteger(); + Runner liveRunner = + Runner.builder() + .app(App.builder().name("test").rootAgent(agent).build()) + .instantSource(() -> Instant.ofEpochMilli(1234L)) + .uuidProvider(() -> String.format("uuid-%04d", counter.getAndIncrement())) + .build(); + Session liveSession = liveRunner.sessionService().createSession("test", "user").blockingGet(); + LiveRequestQueue liveRequestQueue = new LiveRequestQueue(); + TestSubscriber testSubscriber = + liveRunner.runLive(liveSession, liveRequestQueue, RunConfig.builder().build()).test(); + + liveRequestQueue.content(createContent("from user")); + liveRequestQueue.close(); + + testSubscriber.await(); + testSubscriber.assertComplete(); + assertThat(simplifyEvents(testSubscriber.values())).containsExactly("test agent: from llm"); + assertThat(liveSession.id()).isEqualTo("uuid-0000"); + for (Event event : testSubscriber.values()) { + assertThat(event.id()).startsWith("uuid-"); + assertThat(event.invocationId()).startsWith("e-uuid-"); + assertThat(event.timestamp()).isEqualTo(1234L); + } + } + + // A clock fixed at the epoch exposes every event a run builds or rebuilds without the + // invocation's providers: a wall-clock stamp is never 0, and any id not drawn from the counter + // breaks the "uuid-NNNN" pattern. + + @Test + public void runAsync_toolCallTurn_withEpochClock_everyEventFollowsTheProviders() { + DeterministicRun first = runWithEpochClock(toolCallAgent()); + DeterministicRun second = runWithEpochClock(toolCallAgent()); + + // Two parallel calls, their merged responses, and the final text. + assertThat(first.emitted()).hasSize(3); + assertThat(first.emitted().get(0).functionCalls()).hasSize(2); + assertThat(first.emitted().get(1).functionResponses()).hasSize(2); + assertThat(first.emitted().get(2).content().get().text()).isEqualTo("done"); + assertThat(first.persisted()).hasSize(4); + assertEventsFollowProviders(first); + assertIdenticalRuns(first, second); + } + + @Test + public void runAsync_agentCallbacks_withEpochClock_everyEventFollowsTheProviders() { + DeterministicRun first = runWithEpochClock(callbackAgent()); + DeterministicRun second = runWithEpochClock(callbackAgent()); + + // The before-callback's state delta, the model response, and the after-callback's content. + assertThat(first.emitted()).hasSize(3); + assertThat(first.emitted().get(0).actions().stateDelta()).containsEntry("before_key", "v"); + assertThat(first.emitted().get(1).content().get().text()).isEqualTo("from llm"); + assertThat(first.emitted().get(2).content().get().text()).isEqualTo("from callback"); + assertThat(first.persisted()).hasSize(4); + assertEventsFollowProviders(first); + assertIdenticalRuns(first, second); + } + + @Test + public void runAsync_outputSchemaAgent_withEpochClock_everyEventFollowsTheProviders() { + DeterministicRun first = runWithEpochClock(outputSchemaAgent()); + DeterministicRun second = runWithEpochClock(outputSchemaAgent()); + + // The set_model_response call, its response, and the final structured response. + assertThat(first.emitted()).hasSize(3); + assertThat(first.emitted().get(0).functionCalls().get(0).name()) + .hasValue(SetModelResponseTool.NAME); + assertThat(first.emitted().get(2).content().get().text()).isEqualTo("{\"field1\":\"value1\"}"); + assertThat(first.persisted()).hasSize(4); + assertEventsFollowProviders(first); + assertIdenticalRuns(first, second); + } + + /** The events one run emitted and the events it left in the session. */ + private record DeterministicRun(ImmutableList emitted, ImmutableList persisted) {} + + /** Runs {@code agent} once on a fresh runner with an epoch clock and counter-minted ids. */ + private DeterministicRun runWithEpochClock(LlmAgent agent) { + AtomicInteger counter = new AtomicInteger(); + Runner epochRunner = + Runner.builder() + .app(App.builder().name("test").rootAgent(agent).build()) + .instantSource(InstantSource.fixed(Instant.EPOCH)) + .uuidProvider(() -> String.format("uuid-%04d", counter.getAndIncrement())) + .build(); + Session epochSession = + epochRunner + .sessionService() + .createSession("test", "user", null, "fixed-session") + .blockingGet(); + ImmutableList emitted = + ImmutableList.copyOf( + epochRunner + .runAsync("user", epochSession.id(), createContent("hi")) + .toList() + .blockingGet()); + ImmutableList persisted = + ImmutableList.copyOf( + epochRunner + .sessionService() + .getSession("test", "user", epochSession.id(), Optional.empty()) + .blockingGet() + .events()); + return new DeterministicRun(emitted, persisted); + } + + /** An agent whose model issues two id-less {@code echo_tool} calls and then answers in text. */ + private LlmAgent toolCallAgent() { + Part idlessCall = + Part.builder() + .functionCall( + FunctionCall.builder() + .name(echoTool.name()) + .args(ImmutableMap.of("args_name", "args_value")) + .build()) + .build(); + Content parallelCalls = + Content.builder().role("model").parts(ImmutableList.of(idlessCall, idlessCall)).build(); + return createTestAgentBuilder( + createTestLlm(createLlmResponse(parallelCalls), createTextLlmResponse("done"))) + .tools(echoTool) + .build(); + } + + /** + * An agent whose before-agent callback only writes state and whose after-agent callback replies. + */ + private LlmAgent callbackAgent() { + Callbacks.BeforeAgentCallback writeState = + callbackContext -> { + var unused = callbackContext.state().put("before_key", "v"); + return Maybe.empty(); + }; + Callbacks.AfterAgentCallback reply = + callbackContext -> Maybe.just(createContent("from callback")); + return createTestAgentBuilder(createTestLlm(createTextLlmResponse("from llm"))) + .beforeAgentCallback(writeState) + .afterAgentCallback(reply) + .build(); + } + + /** + * An agent with an output schema and a tool under a gemini-2 model name, so the flow routes the + * answer through {@code set_model_response} and builds the final structured response event. + */ + private LlmAgent outputSchemaAgent() { + TestLlm scriptedLlm = + createTestLlm( + createLlmResponse( + Content.fromParts( + Part.fromFunctionCall( + SetModelResponseTool.NAME, ImmutableMap.of("field1", "value1"))))); + BaseLlm gemini2NamedLlm = + new BaseLlm("gemini-2.0-flash") { + @Override + public Flowable generateContent(LlmRequest llmRequest, boolean stream) { + return scriptedLlm.generateContent(llmRequest, stream); + } + + @Override + public BaseLlmConnection connect(LlmRequest llmRequest) { + return scriptedLlm.connect(llmRequest); + } + }; + Schema outputSchema = + Schema.builder() + .type("OBJECT") + .properties(ImmutableMap.of("field1", Schema.builder().type("STRING").build())) + .required(ImmutableList.of("field1")) + .build(); + return LlmAgent.builder() + .name("test agent") + .model(gemini2NamedLlm) + .outputSchema(outputSchema) + .tools(echoTool) + .build(); + } + + /** + * Asserts that every emitted and persisted event carries the epoch timestamp, a counter id, a + * counter-derived invocation id, and ADK-prefixed function-call ids, with no id used twice. + */ + private static void assertEventsFollowProviders(DeterministicRun run) { + for (Event event : Iterables.concat(run.emitted(), run.persisted())) { + assertWithMessage("timestamp of event %s", event.id()).that(event.timestamp()).isEqualTo(0L); + assertThat(event.id()).matches("uuid-\\d{4}"); + assertThat(event.invocationId()).matches("e-uuid-\\d{4}"); + for (FunctionCall functionCall : event.functionCalls()) { + assertThat(functionCall.id().get()).startsWith("adk-uuid-"); + } + for (FunctionResponse functionResponse : event.functionResponses()) { + assertThat(functionResponse.id().get()).startsWith("adk-uuid-"); + } + } + assertThat(ids(run.emitted())).containsNoDuplicates(); + assertThat(ids(run.persisted())).containsNoDuplicates(); + } + + /** Asserts that {@code second} reproduced the ids, timestamps, and call ids of {@code first}. */ + private static void assertIdenticalRuns(DeterministicRun first, DeterministicRun second) { + assertThat(ids(second.emitted())).isEqualTo(ids(first.emitted())); + assertThat(ids(second.persisted())).isEqualTo(ids(first.persisted())); + assertThat(second.persisted().stream().map(Event::timestamp).collect(toImmutableList())) + .isEqualTo(first.persisted().stream().map(Event::timestamp).collect(toImmutableList())); + assertThat(functionCallIds(second.persisted())).isEqualTo(functionCallIds(first.persisted())); + } + + private static ImmutableList ids(List events) { + return events.stream().map(Event::id).collect(toImmutableList()); + } + + private static ImmutableList functionCallIds(List events) { + return events.stream() + .flatMap(event -> event.functionCalls().stream()) + .map(functionCall -> functionCall.id().get()) + .collect(toImmutableList()); + } + + private List runWithDeterministicProviders() { + InstantSource 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()) + .instantSource(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(); + } } 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 2f3a5bcae..14a3e930d 100644 --- a/core/src/test/java/com/google/adk/sessions/InMemorySessionServiceTest.java +++ b/core/src/test/java/com/google/adk/sessions/InMemorySessionServiceTest.java @@ -69,6 +69,23 @@ public void lifecycle_createSession() { assertThat(session.state()).isEmpty(); } + @Test + public void createSession_withDuplicateBehaviorAndProviders_usesBoth() { + InMemorySessionService sessionService = + new InMemorySessionService( + DuplicateSessionIdBehavior.REJECT, + () -> Instant.ofEpochMilli(1234L), + () -> "fixed-session-id"); + + Session session = sessionService.createSession("app-name", "user-id").blockingGet(); + Single duplicate = sessionService.createSession("app-name", "user-id"); + + assertThat(session.id()).isEqualTo("fixed-session-id"); + assertThat(session.lastUpdateTime()).isEqualTo(Instant.ofEpochMilli(1234L)); + SessionException exception = assertThrows(SessionException.class, duplicate::blockingGet); + assertThat(exception).hasMessageThat().isEqualTo(SessionException.SESSION_ALREADY_EXISTS); + } + @Test public void lifecycle_getSession() { InMemorySessionService sessionService = new InMemorySessionService(); diff --git a/core/src/test/java/com/google/adk/summarizer/LlmEventSummarizerTest.java b/core/src/test/java/com/google/adk/summarizer/LlmEventSummarizerTest.java index 37371e987..39b200b24 100644 --- a/core/src/test/java/com/google/adk/summarizer/LlmEventSummarizerTest.java +++ b/core/src/test/java/com/google/adk/summarizer/LlmEventSummarizerTest.java @@ -34,7 +34,9 @@ import com.google.genai.types.FunctionResponse; import com.google.genai.types.Part; import io.reactivex.rxjava3.core.Flowable; +import java.time.Instant; import java.util.Map; +import java.util.concurrent.atomic.AtomicInteger; import org.junit.Before; import org.junit.Rule; import org.junit.Test; @@ -116,6 +118,38 @@ public void summarizeEvents_success() { assertThat(llmRequest.contents().get(0).parts().get().get(0).text()).hasValue(expectedPrompt); } + @Test + public void summarizeEvents_usesInjectedProviders() { + AtomicInteger counter = new AtomicInteger(); + LlmEventSummarizer providerSummarizer = + new LlmEventSummarizer( + mockLlm, + /* promptTemplate= */ null, + () -> Instant.ofEpochMilli(1234L), + () -> String.format("uuid-%04d", counter.getAndIncrement())); + ImmutableList events = + ImmutableList.of(createEvent(1L, "Hello", "user"), createEvent(2L, "Hi there!", "model")); + LlmResponse mockLlmResponse = + LlmResponse.builder() + .content(Content.builder().parts(ImmutableList.of(Part.fromText("Summary"))).build()) + .build(); + when(mockLlm.generateContent(any(LlmRequest.class), eq(false))) + .thenReturn(Flowable.just(mockLlmResponse)); + + Event compactedEvent = providerSummarizer.summarizeEvents(events).blockingGet(); + + // The event id is drawn before the invocation id. + assertThat(compactedEvent.id()).isEqualTo("uuid-0000"); + assertThat(compactedEvent.invocationId()).isEqualTo("uuid-0001"); + assertThat(compactedEvent.timestamp()).isEqualTo(1234L); + // A null prompt template selects the default one. + verify(mockLlm).generateContent(llmRequestCaptor.capture(), eq(false)); + assertThat(llmRequestCaptor.getValue().contents().get(0).parts().get().get(0).text()) + .hasValue( + DEFAULT_PROMPT_TEMPLATE.replace( + "{conversation_history}", "user: Hello\nmodel: Hi there!")); + } + @Test public void summarizeEvents_emptyLlmResponse() { ImmutableList events = ImmutableList.of(createEvent(1L, "Hello", "user")); diff --git a/core/src/test/java/com/google/adk/tools/AgentToolTest.java b/core/src/test/java/com/google/adk/tools/AgentToolTest.java index 113236b63..24db02c45 100644 --- a/core/src/test/java/com/google/adk/tools/AgentToolTest.java +++ b/core/src/test/java/com/google/adk/tools/AgentToolTest.java @@ -23,6 +23,8 @@ import static com.google.adk.testing.TestUtils.createTestLlm; import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertThrows; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; import com.google.adk.agents.BaseAgent; import com.google.adk.agents.Callbacks.AfterAgentCallback; @@ -52,9 +54,11 @@ import io.reactivex.rxjava3.core.Flowable; import io.reactivex.rxjava3.core.Maybe; import io.reactivex.rxjava3.schedulers.Schedulers; +import java.time.Instant; import java.util.Map; import java.util.Optional; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; import org.junit.Before; import org.junit.Test; import org.junit.runner.RunWith; @@ -944,6 +948,77 @@ public void call_nestedRunInheritsCallerScheduler() throws Exception { assertThat(recordingScheduler.workersCreated()).isEqualTo(2); } + @Test + public void call_withMockedInvocationContext_fallsBackToDefaultProvidersAndScheduler() + throws Exception { + LlmAgent testAgent = + createTestAgentBuilder( + createTestLlm( + LlmResponse.builder() + .content(Content.fromParts(Part.fromText("nested response"))) + .build())) + .name("agent_name") + .description("agent description") + .build(); + AgentTool agentTool = AgentTool.create(testAgent); + Session session = + sessionService.createSession("test-app", "test-user", null, "test-session").blockingGet(); + // A Mockito mock answers null from the unstubbed instantSource(), uuidProvider() and + // scheduler() accessors; the nested runner must fall back to the defaults. + InvocationContext mockedContext = mock(InvocationContext.class); + when(mockedContext.agent()).thenReturn(testAgent); + when(mockedContext.session()).thenReturn(session); + when(mockedContext.runConfig()).thenReturn(RunConfig.builder().build()); + ToolContext toolContext = ToolContext.builder(mockedContext).build(); + + Map result = + agentTool.runAsync(ImmutableMap.of("request", "magic"), toolContext).blockingGet(); + + assertThat(result).containsExactly("result", "nested response"); + } + + @Test + public void call_nestedRunner_inheritsParentInvocationProviders() throws Exception { + LlmAgent testAgent = + createTestAgentBuilder( + createTestLlm( + LlmResponse.builder() + .content(Content.fromParts(Part.fromText("nested response"))) + .build())) + .name("agent_name") + .description("agent description") + .build(); + AgentTool agentTool = AgentTool.create(testAgent); + AtomicInteger uuidDraws = new AtomicInteger(); + AtomicInteger clockReads = new AtomicInteger(); + Session session = + sessionService.createSession("test-app", "test-user", null, "test-session").blockingGet(); + ToolContext toolContext = + ToolContext.builder( + InvocationContext.builder() + .invocationId("parent-invocation") + .agent(testAgent) + .session(session) + .sessionService(sessionService) + .instantSource( + () -> { + clockReads.incrementAndGet(); + return Instant.ofEpochMilli(1234L); + }) + .uuidProvider(() -> "uuid-" + uuidDraws.incrementAndGet()) + .build()) + .build(); + + Map result = + agentTool.runAsync(ImmutableMap.of("request", "magic"), toolContext).blockingGet(); + + assertThat(result).containsExactly("result", "nested response"); + // The nested runner's session id, invocation id and event ids/timestamps all came from the + // parent invocation's providers rather than from UUID.randomUUID() / the wall clock. + assertThat(uuidDraws.get()).isAtLeast(3); + assertThat(clockReads.get()).isAtLeast(1); + } + private ToolContext createToolContext(BaseAgent agent) { Session session = sessionService.createSession("test-app", "test-user", null, "test-session").blockingGet();