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();