Skip to content

Commit 742e700

Browse files
DABHDABH
committed
feat: Add injectable time and UUID providers for custom timestamps and IDs
ADK generates timestamps and IDs by calling Instant.now() and UUID.randomUUID() directly, leaving callers no way to control them. This blocks integrations that need to supply their own timestamps and IDs. adk-python and adk-go already expose an equivalent seam. Add a leaf com.google.adk.platform package with TimeProvider and UuidProvider functional interfaces, each with a SYSTEM default that preserves today's wall-clock/random behavior. Rather than an ambient ThreadLocal (which would silently fall back to the system providers once the RxJava flow hops onto a Schedulers worker thread), the providers are threaded as data through InvocationContext, so they are visible on whatever thread builds an event. Callers configure them once on the Runner, which also injects them into the default InMemorySessionService it constructs. Event ids and timestamps, the invocation id, function-call ids, the event compaction summarizer, and the InMemorySessionService session id and lastUpdateTime now derive from the in-scope providers. Event.generateEventId() and the Event.Builder timestamp default delegate to the SYSTEM providers, so the platform interfaces are the single source for generated ids and times while the public Event API stays unchanged.
1 parent 3abcf4f commit 742e700

20 files changed

Lines changed: 643 additions & 45 deletions

‎core/src/main/java/com/google/adk/agents/BaseAgent.java‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -418,7 +418,8 @@ private Maybe<Event> callCallback(
418418
content -> {
419419
invocationContext.setEndInvocation(true);
420420
return Event.builder()
421-
.id(Event.generateEventId())
421+
.id(invocationContext.newUuid())
422+
.timestamp(invocationContext.now().toEpochMilli())
422423
.invocationId(invocationContext.invocationId())
423424
.author(name())
424425
.branch(invocationContext.branch().orElse(null))
@@ -435,7 +436,8 @@ private Maybe<Event> callCallback(
435436
if (callbackContext.state().hasDelta()) {
436437
Event.Builder eventBuilder =
437438
Event.builder()
438-
.id(Event.generateEventId())
439+
.id(invocationContext.newUuid())
440+
.timestamp(invocationContext.now().toEpochMilli())
439441
.invocationId(invocationContext.invocationId())
440442
.author(name())
441443
.branch(invocationContext.branch().orElse(null))

‎core/src/main/java/com/google/adk/agents/InvocationContext.java‎

Lines changed: 67 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,17 +21,19 @@
2121
import com.google.adk.artifacts.BaseArtifactService;
2222
import com.google.adk.memory.BaseMemoryService;
2323
import com.google.adk.models.LlmCallsLimitExceededException;
24+
import com.google.adk.platform.TimeProvider;
25+
import com.google.adk.platform.UuidProvider;
2426
import com.google.adk.plugins.Plugin;
2527
import com.google.adk.plugins.PluginManager;
2628
import com.google.adk.sessions.BaseSessionService;
2729
import com.google.adk.sessions.Session;
2830
import com.google.adk.summarizer.EventsCompactionConfig;
2931
import com.google.errorprone.annotations.CanIgnoreReturnValue;
3032
import com.google.genai.types.Content;
33+
import java.time.Instant;
3134
import java.util.Map;
3235
import java.util.Objects;
3336
import java.util.Optional;
34-
import java.util.UUID;
3537
import java.util.concurrent.ConcurrentHashMap;
3638
import org.jspecify.annotations.Nullable;
3739

@@ -52,6 +54,8 @@ public class InvocationContext {
5254
@Nullable private final ContextCacheConfig contextCacheConfig;
5355
private final InvocationCostManager invocationCostManager;
5456
private final Map<String, Object> callbackContextData;
57+
private final TimeProvider timeProvider;
58+
private final UuidProvider uuidProvider;
5559

5660
@Nullable private String branch;
5761
private BaseAgent agent;
@@ -78,6 +82,8 @@ protected InvocationContext(Builder builder) {
7882
// invocation invocation so that Plugins can access the same data it during the invocation
7983
// across all types of callbacks.
8084
this.callbackContextData = builder.callbackContextData;
85+
this.timeProvider = builder.timeProvider;
86+
this.uuidProvider = builder.uuidProvider;
8187
}
8288

8389
/** Returns a new {@link Builder} for creating {@link InvocationContext} instances. */
@@ -192,9 +198,34 @@ public String userId() {
192198
return session.userId();
193199
}
194200

201+
/** Returns the {@link TimeProvider} for this invocation. */
202+
public TimeProvider timeProvider() {
203+
return timeProvider;
204+
}
205+
206+
/** Returns the {@link UuidProvider} for this invocation. */
207+
public UuidProvider uuidProvider() {
208+
return uuidProvider;
209+
}
210+
211+
/** Returns the current time from this invocation's {@link TimeProvider}. */
212+
public Instant now() {
213+
return timeProvider.now();
214+
}
215+
216+
/** Returns a new unique identifier from this invocation's {@link UuidProvider}. */
217+
public String newUuid() {
218+
return uuidProvider.newUuid();
219+
}
220+
195221
/** Generates a new unique ID for an invocation context. */
196222
public static String newInvocationContextId() {
197-
return "e-" + UUID.randomUUID();
223+
return newInvocationContextId(UuidProvider.SYSTEM);
224+
}
225+
226+
/** Generates a new unique ID for an invocation context using the given {@link UuidProvider}. */
227+
public static String newInvocationContextId(UuidProvider uuidProvider) {
228+
return "e-" + uuidProvider.newUuid();
198229
}
199230

200231
/**
@@ -275,6 +306,8 @@ private Builder(InvocationContext context) {
275306
// invocation invocation so that Plugins can access the same data it during the invocation
276307
// across all types of callbacks.
277308
this.callbackContextData = context.callbackContextData;
309+
this.timeProvider = context.timeProvider;
310+
this.uuidProvider = context.uuidProvider;
278311
}
279312

280313
private BaseSessionService sessionService;
@@ -294,6 +327,8 @@ private Builder(InvocationContext context) {
294327
@Nullable private ContextCacheConfig contextCacheConfig;
295328
private InvocationCostManager invocationCostManager = new InvocationCostManager();
296329
private Map<String, Object> callbackContextData = new ConcurrentHashMap<>();
330+
private TimeProvider timeProvider = TimeProvider.SYSTEM;
331+
private UuidProvider uuidProvider = UuidProvider.SYSTEM;
297332

298333
/**
299334
* Sets the session service for managing session state.
@@ -475,6 +510,30 @@ public Builder callbackContextData(Map<String, Object> callbackContextData) {
475510
return this;
476511
}
477512

513+
/**
514+
* Sets the time provider for the invocation. Defaults to {@link TimeProvider#SYSTEM}.
515+
*
516+
* @param timeProvider the provider for the current time.
517+
* @return this builder instance for chaining.
518+
*/
519+
@CanIgnoreReturnValue
520+
public Builder timeProvider(TimeProvider timeProvider) {
521+
this.timeProvider = timeProvider;
522+
return this;
523+
}
524+
525+
/**
526+
* Sets the UUID provider for the invocation. Defaults to {@link UuidProvider#SYSTEM}.
527+
*
528+
* @param uuidProvider the provider for new unique identifiers.
529+
* @return this builder instance for chaining.
530+
*/
531+
@CanIgnoreReturnValue
532+
public Builder uuidProvider(UuidProvider uuidProvider) {
533+
this.uuidProvider = uuidProvider;
534+
return this;
535+
}
536+
478537
/**
479538
* Builds the {@link InvocationContext} instance.
480539
*
@@ -531,7 +590,9 @@ public boolean equals(Object o) {
531590
&& Objects.equals(eventsCompactionConfig, that.eventsCompactionConfig)
532591
&& Objects.equals(contextCacheConfig, that.contextCacheConfig)
533592
&& Objects.equals(invocationCostManager, that.invocationCostManager)
534-
&& Objects.equals(callbackContextData, that.callbackContextData);
593+
&& Objects.equals(callbackContextData, that.callbackContextData)
594+
&& Objects.equals(timeProvider, that.timeProvider)
595+
&& Objects.equals(uuidProvider, that.uuidProvider);
535596
}
536597

537598
@Override
@@ -553,6 +614,8 @@ public int hashCode() {
553614
eventsCompactionConfig,
554615
contextCacheConfig,
555616
invocationCostManager,
556-
callbackContextData);
617+
callbackContextData,
618+
timeProvider,
619+
uuidProvider);
557620
}
558621
}

‎core/src/main/java/com/google/adk/events/Event.java‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,8 @@
2323
import com.fasterxml.jackson.annotation.JsonProperty;
2424
import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
2525
import com.google.adk.JsonBaseModel;
26+
import com.google.adk.platform.TimeProvider;
27+
import com.google.adk.platform.UuidProvider;
2628
import com.google.common.collect.ImmutableList;
2729
import com.google.common.collect.Iterables;
2830
import com.google.errorprone.annotations.CanIgnoreReturnValue;
@@ -34,12 +36,10 @@
3436
import com.google.genai.types.GenerateContentResponseUsageMetadata;
3537
import com.google.genai.types.GroundingMetadata;
3638
import com.google.genai.types.Transcription;
37-
import java.time.Instant;
3839
import java.util.List;
3940
import java.util.Objects;
4041
import java.util.Optional;
4142
import java.util.Set;
42-
import java.util.UUID;
4343
import org.jspecify.annotations.Nullable;
4444

4545
// TODO - b/413761119 update Agent.java when resolved.
@@ -74,7 +74,7 @@ public class Event extends JsonBaseModel {
7474
private Event() {}
7575

7676
public static String generateEventId() {
77-
return UUID.randomUUID().toString();
77+
return UuidProvider.SYSTEM.newUuid();
7878
}
7979

8080
/** The event id. */
@@ -587,7 +587,7 @@ public Event build() {
587587
event.setCustomMetadata(customMetadata);
588588
event.setModelVersion(modelVersion);
589589
event.setActions(actions().orElseGet(() -> EventActions.builder().build()));
590-
event.setTimestamp(timestamp().orElseGet(() -> Instant.now().toEpochMilli()));
590+
event.setTimestamp(timestamp().orElseGet(() -> TimeProvider.SYSTEM.now().toEpochMilli()));
591591
event.setInputTranscription(inputTranscription);
592592
event.setOutputTranscription(outputTranscription);
593593
return event;

‎core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java‎

Lines changed: 12 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -438,7 +438,7 @@ private Flowable<Event> runOneStep(Context spanContext, InvocationContext contex
438438

439439
final Event mutableEventTemplate =
440440
Event.builder()
441-
.id(Event.generateEventId())
441+
.id(context.newUuid())
442442
.invocationId(context.invocationId())
443443
.author(context.agent().name())
444444
.branch(context.branch().orElse(null))
@@ -453,15 +453,15 @@ private Flowable<Event> runOneStep(Context spanContext, InvocationContext contex
453453
.doFinally(
454454
() -> {
455455
String oldId = mutableEventTemplate.id();
456-
String newId = Event.generateEventId();
456+
String newId = context.newUuid();
457457
logger.debug("Resetting event ID from {} to {}", oldId, newId);
458458
mutableEventTemplate.setId(newId);
459459
})
460460
.concatMap(
461461
event -> {
462462
// Update event ID for the new resulting events
463463
String oldId = event.id();
464-
String newId = Event.generateEventId();
464+
String newId = context.newUuid();
465465
logger.debug("Resetting event ID from {} to {}", oldId, newId);
466466
event = event.toBuilder().id(newId).build();
467467
Flowable<Event> postProcessedEvents = Flowable.just(event);
@@ -555,7 +555,7 @@ public Flowable<Event> runLive(InvocationContext invocationContext) {
555555
return Flowable.empty();
556556
}
557557

558-
String eventIdForSendData = Event.generateEventId();
558+
String eventIdForSendData = invocationContext.newUuid();
559559
LlmAgent agent = (LlmAgent) invocationContext.agent();
560560
BaseLlm llm =
561561
agent.resolvedModel().model().isPresent()
@@ -647,7 +647,7 @@ public void onError(Throwable e) {
647647
.flatMap(
648648
llmResponse -> {
649649
Event baseEventForThisLlmResponse =
650-
liveEventBuilderTemplate.id(Event.generateEventId()).build();
650+
liveEventBuilderTemplate.id(invocationContext.newUuid()).build();
651651
return postprocess(
652652
invocationContext,
653653
baseEventForThisLlmResponse,
@@ -727,7 +727,7 @@ private Flowable<Event> buildPostprocessingEvents(
727727
}
728728

729729
Event modelResponseEvent =
730-
buildModelResponseEvent(baseEventForLlmResponse, llmRequest, updatedResponse);
730+
buildModelResponseEvent(context, baseEventForLlmResponse, llmRequest, updatedResponse);
731731
if (modelResponseEvent.functionCalls().isEmpty()) {
732732
return processorEvents.concatWith(Flowable.just(modelResponseEvent));
733733
}
@@ -772,9 +772,13 @@ private void traceCallLlm(
772772
}
773773

774774
private Event buildModelResponseEvent(
775-
Event baseEventForLlmResponse, LlmRequest llmRequest, LlmResponse llmResponse) {
775+
InvocationContext context,
776+
Event baseEventForLlmResponse,
777+
LlmRequest llmRequest,
778+
LlmResponse llmResponse) {
776779
Event.Builder eventBuilder =
777780
baseEventForLlmResponse.toBuilder()
781+
.timestamp(context.now().toEpochMilli())
778782
.content(llmResponse.content().orElse(null))
779783
.partial(llmResponse.partial().orElse(null))
780784
.errorCode(llmResponse.errorCode().orElse(null))
@@ -794,7 +798,7 @@ private Event buildModelResponseEvent(
794798
logger.debug("event: {} functionCalls: {}", event, event.functionCalls());
795799

796800
if (!event.functionCalls().isEmpty()) {
797-
Functions.populateClientFunctionCallId(event);
801+
Functions.populateClientFunctionCallId(event, context.uuidProvider());
798802
Set<String> longRunningToolIds =
799803
Functions.getLongRunningFunctionCalls(event.functionCalls(), llmRequest.tools());
800804
logger.debug("longRunningToolIds: {}", longRunningToolIds);

‎core/src/main/java/com/google/adk/flows/llmflows/CodeExecution.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -226,6 +226,8 @@ private static Flowable<Event> runPreProcessor(
226226
llmRequest.contents().add(codeContent);
227227
Event codeEvent =
228228
Event.builder()
229+
.id(invocationContext.newUuid())
230+
.timestamp(invocationContext.now().toEpochMilli())
229231
.invocationId(invocationContext.invocationId())
230232
.author(llmAgent.name())
231233
.content(codeContent)
@@ -307,6 +309,8 @@ private static Flowable<Event> runPostProcessor(
307309

308310
Event codeEvent =
309311
Event.builder()
312+
.id(invocationContext.newUuid())
313+
.timestamp(invocationContext.now().toEpochMilli())
310314
.invocationId(invocationContext.invocationId())
311315
.author(llmAgent.name())
312316
.content(responseContent)
@@ -456,6 +460,8 @@ private static Single<Event> postProcessCodeExecutionResult(
456460
}
457461
eventActionsBuilder.artifactDelta(artifactDelta);
458462
return Event.builder()
463+
.id(invocationContext.newUuid())
464+
.timestamp(invocationContext.now().toEpochMilli())
459465
.invocationId(invocationContext.invocationId())
460466
.author(invocationContext.agent().name())
461467
.content(resultContent)

0 commit comments

Comments
 (0)