Fix Vertex session user IDs

Author: Roland Kakonyi

core/src/main/java/com/google/adk/sessions/VertexAiSessionService.java

@@ -35,6 +35,7 @@
 import java.util.Comparator;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
@@ -148,7 +149,7 @@ private ListSessionsResponse parseListSessionsResponse(
       Session session =
           Session.builder(sessionId)
               .appName(appName)
-              .userId(userId)
+              .userId((String) apiSession.get("userId"))
               .state(
                   apiSession.get("sessionState") == null
                       ? new ConcurrentHashMap<>()
@@ -195,6 +196,16 @@ public Maybe<Session> getSession(
         .getSession(reasoningEngineId, sessionId)
         .flatMap(
             getSessionResponseMap -> {
+              String responseUserId =
+                  Optional.ofNullable(getSessionResponseMap.get("userId"))
+                      .map(JsonNode::asText)
+                      .orElse(null);
+              if (!Objects.equals(responseUserId, userId)) {
+                return Maybe.error(
+                    new IllegalArgumentException(
+                        "Session " + sessionId + " does not belong to user " + userId + "."));
+              }
+
               String sessId =
                   Optional.ofNullable(getSessionResponseMap.get("name"))
                       .map(name -> Iterables.getLast(Splitter.on('/').splitToList(name.asText())))

core/src/test/java/com/google/adk/sessions/VertexAiSessionServiceTest.java

@@ -274,6 +274,33 @@ public void listSessions_success() {
     assertThat(sessionsList).hasSize(2);
     ImmutableList<String> ids = sessionsList.stream().map(Session::id).collect(toImmutableList());
     assertThat(ids).containsExactly("1", "2");
+    ImmutableList<String> userIds =
+        sessionsList.stream().map(Session::userId).collect(toImmutableList());
+    assertThat(userIds).containsExactly("user", "user");
+  }
+
+  @Test
+  public void listSessions_usesResponseUserId() throws Exception {
+    when(mockApiClient.request("GET", "reasoningEngines/123/sessions?filter=user_id=user1", ""))
+        .thenAnswer(
+            new MockApiAnswer(
+                """
+                {
+                  "sessions": [
+                    {
+                      "name": "projects/test-project/locations/test-location/reasoningEngines/123/sessions/3",
+                      "userId": "user2",
+                      "updateTime": "2024-12-14T12:12:12.123456Z"
+                    }
+                  ]
+                }\
+                """));
+
+    ListSessionsResponse sessions =
+        vertexAiSessionService.listSessions("123", "user1").blockingGet();
+
+    assertThat(sessions.sessions()).hasSize(1);
+    assertThat(sessions.sessions().get(0).userId()).isEqualTo("user2");
   }
 
   @Test
@@ -346,12 +373,25 @@ public void listEvents_empty() {
   public void listEmptySession_success() {
     assertThat(
             vertexAiSessionService
-                .getSession("789", "user1", "3", Optional.empty())
+                .getSession("789", "user2", "3", Optional.empty())
                 .blockingGet()
                 .events())
         .isEmpty();
   }
 
+  @Test
+  public void getSession_whenResponseUserIdDiffers_throws() {
+    IllegalArgumentException exception =
+        assertThrows(
+            IllegalArgumentException.class,
+            () ->
+                vertexAiSessionService
+                    .getSession("789", "user1", "3", Optional.empty())
+                    .blockingGet());
+
+    assertThat(exception).hasMessageThat().contains("Session 3 does not belong to user user1.");
+  }
+
   @Test
   public void appendEvent_withStateRemoved_updatesSessionState() {
     String userId = "userB";