Skip to content

Commit 2daed36

Browse files
google-genai-botcopybara-github
authored andcommitted
fix: Bypass redundant getSession read in ADK Runner
Bypass the sequential getSession database call after appending a message in Runner.runAsyncImpl to save one network roundtrip (~170ms). PiperOrigin-RevId: 927833117
1 parent ec93f50 commit 2daed36

2 files changed

Lines changed: 42 additions & 13 deletions

File tree

core/src/main/java/com/google/adk/runner/Runner.java

Lines changed: 3 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -518,18 +518,9 @@ protected Flowable<Event> runAsyncImpl(
518518
runConfig.saveInputBlobsAsArtifacts(),
519519
stateDelta))
520520
.flatMapPublisher(
521-
event -> {
522-
// Get the updated session after the message and state delta are
523-
// applied
524-
return this.sessionService
525-
.getSession(
526-
session.appName(), session.userId(), session.id(), Optional.empty())
527-
.flatMapPublisher(
528-
updatedSession ->
529-
runAgentWithUpdatedSession(
530-
initialContext, updatedSession, event, rootAgent))
531-
.compose(Tracing.<Event>withContext(capturedContext));
532-
});
521+
event ->
522+
runAgentWithUpdatedSession(initialContext, session, event, rootAgent)
523+
.compose(Tracing.<Event>withContext(capturedContext)));
533524
})
534525
.doOnError(
535526
throwable -> {

core/src/test/java/com/google/adk/runner/RunnerTest.java

Lines changed: 39 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -763,8 +763,11 @@ public void runAsync_concurrentCalls_staleRead() throws Exception {
763763
.firstOrError()
764764
.doOnSuccess(
765765
event -> {
766+
s.events().add(e);
767+
if (e.actions() != null && e.actions().stateDelta() != null) {
768+
s.state().putAll(e.actions().stateDelta());
769+
}
766770
List<Event> newEvents = new ArrayList<>(s.events());
767-
newEvents.add(e);
768771
Session updated =
769772
Session.builder(s.id())
770773
.appName(s.appName())
@@ -1193,6 +1196,41 @@ public void runAsync_pureMockSessionService_multiStepLlmAgent_appendsExactlyOnce
11931196
verify(pureMockSessionService, times(expectedAppendCount)).appendEvent(any(), any());
11941197
}
11951198

1199+
@Test
1200+
public void runAsync_bypassesRedundantGetSession() {
1201+
BaseSessionService mockSessionService = mock(BaseSessionService.class);
1202+
Session backingSession = Session.builder("session-id").appName("test").userId("user").build();
1203+
1204+
when(mockSessionService.getSession(anyString(), anyString(), anyString(), any()))
1205+
.thenReturn(Maybe.just(backingSession));
1206+
when(mockSessionService.appendEvent(any(), any()))
1207+
.thenReturn(Single.just(Event.builder().id("sentinel").author("user").build()));
1208+
1209+
BaseAgent mockAgent = mock(BaseAgent.class);
1210+
when(mockAgent.runAsync(any()))
1211+
.thenReturn(Flowable.just(Event.builder().id("agent-event").author("agent").build()));
1212+
1213+
Runner spyRunner =
1214+
Runner.builder()
1215+
.app(
1216+
App.builder()
1217+
.name("test")
1218+
.rootAgent(mockAgent)
1219+
.plugins(ImmutableList.of(plugin))
1220+
.build())
1221+
.sessionService(mockSessionService)
1222+
.build();
1223+
1224+
List<Event> unused =
1225+
spyRunner
1226+
.runAsync("user", backingSession.id(), createContent("from user"))
1227+
.toList()
1228+
.blockingGet();
1229+
1230+
// Verify getSession was only called once (at the start of runAsync)
1231+
verify(mockSessionService, times(1)).getSession(anyString(), anyString(), anyString(), any());
1232+
}
1233+
11961234
@Test
11971235
public void runAsync_withSessionKey_success() {
11981236
var events =

0 commit comments

Comments
 (0)