From 77eae25844dbc7d49101461c97d7a8711ba5e9dd Mon Sep 17 00:00:00 2001 From: Kamil Tomaszek Date: Wed, 9 Sep 2026 14:36:34 -0700 Subject: [PATCH] fix(sessions): forward the caller-supplied session id in VertexAiSessionService `createSession` accepted a session id but never sent it, so the backend always minted its own. It now travels as the `sessionId` query parameter and is validated first, so an id that fails validation, or that the backend rejects, now errors instead of quietly succeeding under a different id. PiperOrigin-RevId: 978758795 --- .../google/adk/sessions/VertexAiClient.java | 22 +++-- .../adk/sessions/VertexAiSessionService.java | 24 ++++-- .../google/adk/sessions/MockApiAnswer.java | 47 ++++++++--- .../sessions/VertexAiSessionServiceTest.java | 80 +++++++++++++++++++ 4 files changed, 152 insertions(+), 21 deletions(-) diff --git a/core/src/main/java/com/google/adk/sessions/VertexAiClient.java b/core/src/main/java/com/google/adk/sessions/VertexAiClient.java index d6e74fe1c..310c8f93f 100644 --- a/core/src/main/java/com/google/adk/sessions/VertexAiClient.java +++ b/core/src/main/java/com/google/adk/sessions/VertexAiClient.java @@ -64,19 +64,30 @@ final class VertexAiClient { this.apiClient = new HttpApiClient(project, location, credentials, httpOptions); } + /** + * Creates a session, optionally under a caller-chosen {@code sessionId}. The backend mints an id + * when none is supplied. + */ Maybe createSession( - String reasoningEngineId, String userId, Map state) { + String reasoningEngineId, + String userId, + @Nullable Map state, + @Nullable String sessionId) { Map sessionJsonMap = new HashMap<>(); sessionJsonMap.put("userId", userId); if (state != null) { sessionJsonMap.put("sessionState", state); } + String createPath = + "reasoningEngines/" + + reasoningEngineId + + "/sessions" + + (sessionId == null + ? "" + : "?sessionId=" + UrlEscapers.urlFormParameterEscaper().escape(sessionId)); return Single.fromCallable(() -> objectMapper.writeValueAsString(sessionJsonMap)) - .flatMap( - sessionJson -> - performApiRequest( - "POST", "reasoningEngines/" + reasoningEngineId + "/sessions", sessionJson)) + .flatMap(sessionJson -> performApiRequest("POST", createPath, sessionJson)) .flatMapMaybe( apiResponse -> { logger.debug("Create Session response {}", apiResponse.getResponseBody()); @@ -86,6 +97,7 @@ Maybe createSession( jsonResponse -> { String sessionName = jsonResponse.get("name").asText(); List parts = Splitter.on('/').splitToList(sessionName); + // The backend is authoritative and mints its own id when none was supplied. String sessId = parts.get(parts.size() - 3); String operationId = Iterables.getLast(parts); diff --git a/core/src/main/java/com/google/adk/sessions/VertexAiSessionService.java b/core/src/main/java/com/google/adk/sessions/VertexAiSessionService.java index 92c10cd97..8ede5cc0b 100644 --- a/core/src/main/java/com/google/adk/sessions/VertexAiSessionService.java +++ b/core/src/main/java/com/google/adk/sessions/VertexAiSessionService.java @@ -25,6 +25,7 @@ import com.google.adk.events.Event; import com.google.auth.oauth2.GoogleCredentials; import com.google.common.base.Splitter; +import com.google.common.base.Strings; import com.google.common.collect.Iterables; import com.google.genai.types.HttpOptions; import io.reactivex.rxjava3.core.Completable; @@ -79,19 +80,32 @@ public Single createSession( return createSession(appName, userId, (Map) state, sessionId); } + /** + * Creates a session, requesting {@code sessionId} as its id when one is given. A non-empty id + * must match {@code [a-zA-Z0-9_-]+}; an empty or null one asks the backend to generate it. + * + *

The backend's rule is narrower still - up to 63 characters from {@code [a-z0-9-]}, starting + * with a letter and ending with a letter or digit - so it rejects some ids that pass validation + * here. + */ @Override public Single createSession( String appName, String userId, @Nullable Map state, @Nullable String sessionId) { + // Empty means "generate one" just as null does, per this method's contract. + String requestedSessionId = Strings.emptyToNull(sessionId); + if (requestedSessionId != null) { + validateSessionId(requestedSessionId); + } String reasoningEngineId = parseReasoningEngineId(appName); return client - .createSession(reasoningEngineId, userId, state) + .createSession(reasoningEngineId, userId, state, requestedSessionId) .map( getSessionResponseMap -> - parseSession(getSessionResponseMap, appName, userId, sessionId)) + parseSession(getSessionResponseMap, appName, userId, requestedSessionId)) .toSingle(); } @@ -125,14 +139,12 @@ public Single listSessions(String appName, String userId) return client .listSessions(reasoningEngineId, userId) - .map( - listSessionsResponseMap -> - parseListSessionsResponse(listSessionsResponseMap, appName, userId)) + .map(listSessionsResponseMap -> parseListSessionsResponse(listSessionsResponseMap, appName)) .defaultIfEmpty(ListSessionsResponse.builder().sessions(new ArrayList<>()).build()); } private ListSessionsResponse parseListSessionsResponse( - JsonNode listSessionsResponseMap, String appName, String userId) { + JsonNode listSessionsResponseMap, String appName) { JsonNode sessionsNode = listSessionsResponseMap.get("sessions"); if (sessionsNode == null || sessionsNode.isNull() || sessionsNode.isEmpty()) { return ListSessionsResponse.builder().build(); diff --git a/core/src/test/java/com/google/adk/sessions/MockApiAnswer.java b/core/src/test/java/com/google/adk/sessions/MockApiAnswer.java index 84c860996..864ca5525 100644 --- a/core/src/test/java/com/google/adk/sessions/MockApiAnswer.java +++ b/core/src/test/java/com/google/adk/sessions/MockApiAnswer.java @@ -16,13 +16,14 @@ package com.google.adk.sessions; +import static java.nio.charset.StandardCharsets.UTF_8; + import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.ObjectMapper; import com.google.adk.JsonBaseModel; import com.google.adk.events.Event; import java.io.IOException; import java.net.URLDecoder; -import java.nio.charset.StandardCharsets; import java.time.Instant; import java.util.ArrayList; import java.util.HashMap; @@ -44,7 +45,7 @@ class MockApiAnswer implements Answer { private static final Pattern SESSION_REGEX = Pattern.compile("^reasoningEngines/([^/]+)/sessions/([^/]+)$"); private static final Pattern SESSIONS_REGEX = - Pattern.compile("^reasoningEngines/([^/]+)/sessions$"); + Pattern.compile("^reasoningEngines/([^/]+)/sessions(?:\\?sessionId=(.+))?$"); private static final Pattern SESSIONS_FILTER_REGEX = Pattern.compile("^reasoningEngines/([^/]+)/sessions\\?filter=(.+)$"); private static final String USER_ID_FILTER_PREFIX = "user_id="; @@ -56,6 +57,17 @@ class MockApiAnswer implements Answer { private static final MediaType JSON_MEDIA_TYPE = MediaType.parse("application/json; charset=utf-8"); + /** The id this fake mints when the caller requests none. */ + static final String GENERATED_SESSION_ID = "4"; + + /** + * A requested id this fake deliberately overrides, so a test can assert the backend's id wins. + */ + static final String OVERRIDDEN_REQUEST_ID = "backend-renames-this"; + + /** The id this fake substitutes for {@link #OVERRIDDEN_REQUEST_ID}. */ + static final String BACKEND_CHOSEN_ID = "backend-chosen-id"; + private final Map sessionMap; private final Map eventMap; private final String rawApiResponse; @@ -121,12 +133,30 @@ public void close() {} private ApiResponse handleCreateSession(String path, InvocationOnMock invocation) throws Exception { - String newSessionId = "4"; + Matcher sessionsMatcher = SESSIONS_REGEX.matcher(path); + if (!sessionsMatcher.matches()) { + return null; + } + // Create the session under the caller-supplied id, as the real service does. + String basePath = "reasoningEngines/" + sessionsMatcher.group(1) + "/sessions"; + String requestedSessionId = + sessionsMatcher.group(2) == null + ? null + : URLDecoder.decode(sessionsMatcher.group(2), UTF_8); + String newSessionId; + if (requestedSessionId == null) { + newSessionId = GENERATED_SESSION_ID; + } else if (requestedSessionId.equals(OVERRIDDEN_REQUEST_ID)) { + // Lets a test prove the response id wins over the requested one. + newSessionId = BACKEND_CHOSEN_ID; + } else { + newSessionId = requestedSessionId; + } Map requestDict = mapper.readValue( (String) invocation.getArgument(2), new TypeReference>() {}); Map newSessionData = new HashMap<>(); - newSessionData.put("name", path + "/" + newSessionId); + newSessionData.put("name", basePath + "/" + newSessionId); newSessionData.put("userId", requestDict.get("userId")); newSessionData.put("sessionState", requestDict.get("sessionState")); newSessionData.put("updateTime", "2024-12-12T12:12:12.123456Z"); @@ -141,7 +171,7 @@ private ApiResponse handleCreateSession(String path, InvocationOnMock invocation "done": false } """, - path, newSessionId)); + basePath, newSessionId)); } private ApiResponse handleGetSession(String path) throws Exception { @@ -165,7 +195,7 @@ private ApiResponse handleGetSessions(String path) throws Exception { // Decode the URL-escaped filter and read the quoted user_id literal back with // a JSON parser, as the real server would. An unquoted/injected filter is // rejected. - String decodedFilter = URLDecoder.decode(sessionsMatcher.group(2), StandardCharsets.UTF_8); + String decodedFilter = URLDecoder.decode(sessionsMatcher.group(2), UTF_8); if (!decodedFilter.startsWith(USER_ID_FILTER_PREFIX)) { throw new IllegalArgumentException("Unsupported sessions filter: " + decodedFilter); } @@ -245,10 +275,7 @@ private ApiResponse handleGetEvents(String path) throws Exception { } String sessionId = matcher.group(2); // The client URL-escapes the filter value; decode it as the real server would. - String filter = - matcher.group(3) == null - ? null - : URLDecoder.decode(matcher.group(3), StandardCharsets.UTF_8); + String filter = matcher.group(3) == null ? null : URLDecoder.decode(matcher.group(3), UTF_8); String eventData = eventMap.get(sessionId); if (eventData != null) { if (filter != null) { diff --git a/core/src/test/java/com/google/adk/sessions/VertexAiSessionServiceTest.java b/core/src/test/java/com/google/adk/sessions/VertexAiSessionServiceTest.java index bd703314e..1cf7c247d 100644 --- a/core/src/test/java/com/google/adk/sessions/VertexAiSessionServiceTest.java +++ b/core/src/test/java/com/google/adk/sessions/VertexAiSessionServiceTest.java @@ -22,6 +22,7 @@ import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.atLeastOnce; +import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -207,6 +208,85 @@ public void createSession_success() throws Exception { assertThat(newSessionMap.get("sessionState")).isEqualTo(sessionStateMap); } + @Test + public void createSession_sessionId_createsSessionUnderThatId() throws Exception { + Session createdSession = + vertexAiSessionService + .createSession("123", "test_user", (Map) null, "my-session") + .blockingGet(); + + // The id reaches the fake only through the request query string. + assertThat(createdSession.id()).isEqualTo("my-session"); + assertThat(sessionMap).containsKey("my-session"); + } + + @Test + public void createSession_backendChoosesDifferentId_returnsBackendId() throws Exception { + Session createdSession = + vertexAiSessionService + .createSession( + "123", "test_user", (Map) null, MockApiAnswer.OVERRIDDEN_REQUEST_ID) + .blockingGet(); + + // The backend is authoritative: its id wins over the one the caller asked for. + assertThat(createdSession.id()).isEqualTo(MockApiAnswer.BACKEND_CHOSEN_ID); + } + + @Test + public void createSession_sessionIdWithReservedCharacters_escapedIntoQuery() throws Exception { + // VertexAiClient is reachable without the service's validation, so it must escape the id. + ArgumentCaptor path = ArgumentCaptor.forClass(String.class); + Object unused = + new VertexAiClient("test-project", "test-location", mockApiClient) + .createSession("123", "user", null, "a b&c=d") + .blockingGet(); + + verify(mockApiClient).request(eq("POST"), path.capture(), anyString()); + assertThat(path.getValue()).isEqualTo("reasoningEngines/123/sessions?sessionId=a+b%26c%3Dd"); + } + + @Test + public void createSession_emptySessionId_usesBackendGeneratedId() throws Exception { + // The interface contract treats an empty id the same way as a null one. + Session createdSession = + vertexAiSessionService + .createSession("123", "test_user", (Map) null, "") + .blockingGet(); + + assertThat(createdSession.id()).isEqualTo(MockApiAnswer.GENERATED_SESSION_ID); + } + + @Test + public void createSession_invalidSessionId_throwsWithoutCallingBackend() throws Exception { + // Whitespace is not empty, so it reaches the allowlist and is rejected there. + for (String bad : ImmutableList.of("bad/id", " ")) { + assertThrows( + IllegalArgumentException.class, + () -> + vertexAiSessionService.createSession( + "123", "test_user", (Map) null, bad)); + } + + verify(mockApiClient, never()).request(anyString(), anyString(), anyString()); + } + + @Test + public void createSession_backendRejectsSessionId_propagatesAsError() { + // The local allowlist is looser than the backend's, so a locally valid id can still be refused. + when(mockApiClient.request( + eq("POST"), eq("reasoningEngines/123/sessions?sessionId=My_Session_ID"), anyString())) + .thenReturn(MockApiAnswer.responseWithStatus(400, "{\"error\": \"invalid session_id\"}")); + + // The id clears local validation, so nothing throws until the call is subscribed. + Single session = + vertexAiSessionService.createSession( + "123", "test_user", (Map) null, "My_Session_ID"); + + VertexAiApiException exception = + assertThrows(VertexAiApiException.class, () -> session.blockingGet()); + assertThat(exception.statusCode()).isEqualTo(400); + } + @Test public void createSession_getSession_success() throws Exception { Map sessionStateMap = new HashMap<>(ImmutableMap.of("new_key", "new_value"));