Skip to content

Commit 4e43c28

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 4225b07 commit 4e43c28

21 files changed

Lines changed: 648 additions & 46 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
@@ -419,7 +419,8 @@ private Maybe<Event> callCallback(
419419
content -> {
420420
invocationContext.setEndInvocation(true);
421421
return Event.builder()
422-
.id(Event.generateEventId())
422+
.id(invocationContext.newUuid())
423+
.timestamp(invocationContext.now().toEpochMilli())
423424
.invocationId(invocationContext.invocationId())
424425
.author(name())
425426
.branch(invocationContext.branch().orElse(null))
@@ -436,7 +437,8 @@ private Maybe<Event> callCallback(
436437
if (callbackContext.state().hasDelta()) {
437438
Event.Builder eventBuilder =
438439
Event.builder()
439-
.id(Event.generateEventId())
440+
.id(invocationContext.newUuid())
441+
.timestamp(invocationContext.now().toEpochMilli())
440442
.invocationId(invocationContext.invocationId())
441443
.author(name())
442444
.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
@@ -22,17 +22,19 @@
2222
import com.google.adk.artifacts.BaseArtifactService;
2323
import com.google.adk.memory.BaseMemoryService;
2424
import com.google.adk.models.LlmCallsLimitExceededException;
25+
import com.google.adk.platform.TimeProvider;
26+
import com.google.adk.platform.UuidProvider;
2527
import com.google.adk.plugins.Plugin;
2628
import com.google.adk.plugins.PluginManager;
2729
import com.google.adk.sessions.BaseSessionService;
2830
import com.google.adk.sessions.Session;
2931
import com.google.adk.summarizer.EventsCompactionConfig;
3032
import com.google.errorprone.annotations.CanIgnoreReturnValue;
3133
import com.google.genai.types.Content;
34+
import java.time.Instant;
3235
import java.util.Map;
3336
import java.util.Objects;
3437
import java.util.Optional;
35-
import java.util.UUID;
3638
import java.util.concurrent.ConcurrentHashMap;
3739
import org.jspecify.annotations.Nullable;
3840

@@ -55,6 +57,8 @@ public class InvocationContext {
5557
private final @Nullable ResumabilityConfig resumabilityConfig;
5658
private final InvocationCostManager invocationCostManager;
5759
private final Map<String, Object> callbackContextData;
60+
private final TimeProvider timeProvider;
61+
private final UuidProvider uuidProvider;
5862

5963
@Nullable private String branch;
6064
private BaseAgent agent;
@@ -82,6 +86,8 @@ protected InvocationContext(Builder builder) {
8286
// invocation invocation so that Plugins can access the same data it during the invocation
8387
// across all types of callbacks.
8488
this.callbackContextData = builder.callbackContextData;
89+
this.timeProvider = builder.timeProvider;
90+
this.uuidProvider = builder.uuidProvider;
8591
}
8692

8793
/** Returns a new {@link Builder} for creating {@link InvocationContext} instances. */
@@ -196,9 +202,34 @@ public String userId() {
196202
return session.userId();
197203
}
198204

205+
/** Returns the {@link TimeProvider} for this invocation. */
206+
public TimeProvider timeProvider() {
207+
return timeProvider;
208+
}
209+
210+
/** Returns the {@link UuidProvider} for this invocation. */
211+
public UuidProvider uuidProvider() {
212+
return uuidProvider;
213+
}
214+
215+
/** Returns the current time from this invocation's {@link TimeProvider}. */
216+
public Instant now() {
217+
return timeProvider.now();
218+
}
219+
220+
/** Returns a new unique identifier from this invocation's {@link UuidProvider}. */
221+
public String newUuid() {
222+
return uuidProvider.newUuid();
223+
}
224+
199225
/** Generates a new unique ID for an invocation context. */
200226
public static String newInvocationContextId() {
201-
return "e-" + UUID.randomUUID();
227+
return newInvocationContextId(UuidProvider.SYSTEM);
228+
}
229+
230+
/** Generates a new unique ID for an invocation context using the given {@link UuidProvider}. */
231+
public static String newInvocationContextId(UuidProvider uuidProvider) {
232+
return "e-" + uuidProvider.newUuid();
202233
}
203234

204235
/**
@@ -288,6 +319,8 @@ private Builder(InvocationContext context) {
288319
// invocation invocation so that Plugins can access the same data it during the invocation
289320
// across all types of callbacks.
290321
this.callbackContextData = context.callbackContextData;
322+
this.timeProvider = context.timeProvider;
323+
this.uuidProvider = context.uuidProvider;
291324
}
292325

293326
private BaseSessionService sessionService;
@@ -308,6 +341,8 @@ private Builder(InvocationContext context) {
308341
private @Nullable ResumabilityConfig resumabilityConfig;
309342
private InvocationCostManager invocationCostManager = new InvocationCostManager();
310343
private Map<String, Object> callbackContextData = new ConcurrentHashMap<>();
344+
private TimeProvider timeProvider = TimeProvider.SYSTEM;
345+
private UuidProvider uuidProvider = UuidProvider.SYSTEM;
311346

312347
/**
313348
* Sets the session service for managing session state.
@@ -501,6 +536,30 @@ public Builder callbackContextData(Map<String, Object> callbackContextData) {
501536
return this;
502537
}
503538

539+
/**
540+
* Sets the time provider for the invocation. Defaults to {@link TimeProvider#SYSTEM}.
541+
*
542+
* @param timeProvider the provider for the current time.
543+
* @return this builder instance for chaining.
544+
*/
545+
@CanIgnoreReturnValue
546+
public Builder timeProvider(TimeProvider timeProvider) {
547+
this.timeProvider = timeProvider;
548+
return this;
549+
}
550+
551+
/**
552+
* Sets the UUID provider for the invocation. Defaults to {@link UuidProvider#SYSTEM}.
553+
*
554+
* @param uuidProvider the provider for new unique identifiers.
555+
* @return this builder instance for chaining.
556+
*/
557+
@CanIgnoreReturnValue
558+
public Builder uuidProvider(UuidProvider uuidProvider) {
559+
this.uuidProvider = uuidProvider;
560+
return this;
561+
}
562+
504563
/**
505564
* Builds the {@link InvocationContext} instance.
506565
*
@@ -558,7 +617,9 @@ public boolean equals(Object o) {
558617
&& Objects.equals(contextCacheConfig, that.contextCacheConfig)
559618
&& Objects.equals(resumabilityConfig, that.resumabilityConfig)
560619
&& Objects.equals(invocationCostManager, that.invocationCostManager)
561-
&& Objects.equals(callbackContextData, that.callbackContextData);
620+
&& Objects.equals(callbackContextData, that.callbackContextData)
621+
&& Objects.equals(timeProvider, that.timeProvider)
622+
&& Objects.equals(uuidProvider, that.uuidProvider);
562623
}
563624

564625
@Override
@@ -581,6 +642,8 @@ public int hashCode() {
581642
contextCacheConfig,
582643
resumabilityConfig,
583644
invocationCostManager,
584-
callbackContextData);
645+
callbackContextData,
646+
timeProvider,
647+
uuidProvider);
585648
}
586649
}

‎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
@@ -451,7 +451,7 @@ private Flowable<Event> runOneStep(Context spanContext, InvocationContext contex
451451

452452
final Event mutableEventTemplate =
453453
Event.builder()
454-
.id(Event.generateEventId())
454+
.id(context.newUuid())
455455
.invocationId(context.invocationId())
456456
.author(context.agent().name())
457457
.branch(context.branch().orElse(null))
@@ -466,15 +466,15 @@ private Flowable<Event> runOneStep(Context spanContext, InvocationContext contex
466466
.doFinally(
467467
() -> {
468468
String oldId = mutableEventTemplate.id();
469-
String newId = Event.generateEventId();
469+
String newId = context.newUuid();
470470
logger.debug("Resetting event ID from {} to {}", oldId, newId);
471471
mutableEventTemplate.setId(newId);
472472
})
473473
.concatMap(
474474
event -> {
475475
// Update event ID for the new resulting events
476476
String oldId = event.id();
477-
String newId = Event.generateEventId();
477+
String newId = context.newUuid();
478478
logger.debug("Resetting event ID from {} to {}", oldId, newId);
479479
event = event.toBuilder().id(newId).build();
480480
Flowable<Event> postProcessedEvents = Flowable.just(event);
@@ -575,7 +575,7 @@ public Flowable<Event> runLive(InvocationContext invocationContext) {
575575
return Flowable.empty();
576576
}
577577

578-
String eventIdForSendData = Event.generateEventId();
578+
String eventIdForSendData = invocationContext.newUuid();
579579
LlmAgent agent = (LlmAgent) invocationContext.agent();
580580
BaseLlm llm =
581581
agent.resolvedModel().model().isPresent()
@@ -667,7 +667,7 @@ public void onError(Throwable e) {
667667
.flatMap(
668668
llmResponse -> {
669669
Event baseEventForThisLlmResponse =
670-
liveEventBuilderTemplate.id(Event.generateEventId()).build();
670+
liveEventBuilderTemplate.id(invocationContext.newUuid()).build();
671671
return postprocess(
672672
invocationContext,
673673
baseEventForThisLlmResponse,
@@ -750,7 +750,7 @@ private Flowable<Event> buildPostprocessingEvents(
750750
}
751751

752752
Event modelResponseEvent =
753-
buildModelResponseEvent(baseEventForLlmResponse, llmRequest, updatedResponse);
753+
buildModelResponseEvent(context, baseEventForLlmResponse, llmRequest, updatedResponse);
754754
if (modelResponseEvent.functionCalls().isEmpty()
755755
|| modelResponseEvent.partial().orElse(false)) {
756756
return processorEvents.concatWith(Flowable.just(modelResponseEvent));
@@ -796,9 +796,13 @@ private void traceCallLlm(
796796
}
797797

798798
private Event buildModelResponseEvent(
799-
Event baseEventForLlmResponse, LlmRequest llmRequest, LlmResponse llmResponse) {
799+
InvocationContext context,
800+
Event baseEventForLlmResponse,
801+
LlmRequest llmRequest,
802+
LlmResponse llmResponse) {
800803
Event.Builder eventBuilder =
801804
baseEventForLlmResponse.toBuilder()
805+
.timestamp(context.now().toEpochMilli())
802806
.content(llmResponse.content().orElse(null))
803807
.partial(llmResponse.partial().orElse(null))
804808
.errorCode(llmResponse.errorCode().orElse(null))
@@ -818,7 +822,7 @@ private Event buildModelResponseEvent(
818822
logger.debug("event: {} functionCalls: {}", event, event.functionCalls());
819823

820824
if (!event.functionCalls().isEmpty()) {
821-
Functions.populateClientFunctionCallId(event);
825+
Functions.populateClientFunctionCallId(event, context.uuidProvider());
822826
Set<String> longRunningToolIds =
823827
Functions.getLongRunningFunctionCalls(event.functionCalls(), llmRequest.tools());
824828
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)