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 3b8ce74ed..3f2393c36 100644 --- a/core/src/main/java/com/google/adk/sessions/InMemorySessionService.java +++ b/core/src/main/java/com/google/adk/sessions/InMemorySessionService.java @@ -251,6 +251,8 @@ public Single appendEvent(Session session, Event event) { .computeIfAbsent(appName, k -> new ConcurrentHashMap<>()) .computeIfAbsent(userId, k -> new ConcurrentHashMap<>()) .put(userStateKey, value); + } else { + session.state().put(key, value); } }); } @@ -265,6 +267,8 @@ public Single appendEvent(Session session, Event event) { .getOrDefault(userId, new ConcurrentHashMap<>()) .put(sessionId, session); + mergeWithGlobalState(appName, userId, session); + return Single.just(event); } 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 eed693254..d30081c71 100644 --- a/core/src/test/java/com/google/adk/sessions/InMemorySessionServiceTest.java +++ b/core/src/test/java/com/google/adk/sessions/InMemorySessionServiceTest.java @@ -107,32 +107,36 @@ public void lifecycle_deleteSession() { } @Test - public void appendEvent_withStateDelta_mutatesStateCorrectly() { - ConcurrentMap startingState1 = new ConcurrentHashMap<>(); - startingState1.put("key_to_be_removed", "value_to_be_removed"); - ConcurrentMap startingState2 = new ConcurrentHashMap<>(); - startingState2.put("key_to_be_removed", "value_to_be_removed"); - - ConcurrentMap stateDelta = new ConcurrentHashMap<>(); - State expectedFinalState = new State(startingState1, stateDelta); - expectedFinalState.remove("key_to_be_removed"); - expectedFinalState.put("new_key", "new_value"); - + public void appendEvent_updatesSessionState() { InMemorySessionService sessionService = new InMemorySessionService(); Session session = sessionService - .createSession("app-name", "user-id", startingState2, /* sessionId= */ null) + .createSession("app", "user", new ConcurrentHashMap<>(), "session1") .blockingGet(); - var unused = + ConcurrentMap stateDelta = new ConcurrentHashMap<>(); + stateDelta.put("sessionKey", "sessionValue"); + stateDelta.put("_app_appKey", "appValue"); + stateDelta.put("_user_userKey", "userValue"); + + Event event = + Event.builder().actions(EventActions.builder().stateDelta(stateDelta).build()).build(); + + var unused = sessionService.appendEvent(session, event).blockingGet(); + + // After appendEvent, session state in memory should contain session-specific state from delta + // and merged global state. + assertThat(session.state()).containsEntry("sessionKey", "sessionValue"); + assertThat(session.state()).containsEntry("_app_appKey", "appValue"); + assertThat(session.state()).containsEntry("_user_userKey", "userValue"); + + // getSession should return session with merged state. + Session retrievedSession = sessionService - .appendEvent( - session, - Event.builder() - .actions(EventActions.builder().stateDelta(stateDelta).build()) - .build()) + .getSession(session.appName(), session.userId(), session.id(), Optional.empty()) .blockingGet(); - - assertThat(session.state().entrySet()).containsExactlyElementsIn(expectedFinalState.entrySet()); + assertThat(retrievedSession.state()).containsEntry("sessionKey", "sessionValue"); + assertThat(retrievedSession.state()).containsEntry("_app_appKey", "appValue"); + assertThat(retrievedSession.state()).containsEntry("_user_userKey", "userValue"); } }