From a3ca97584bb32e5e23acb601d4e6471b8e32d6e1 Mon Sep 17 00:00:00 2001 From: AnupamKumar-1 Date: Sat, 29 Aug 2026 01:35:21 +0530 Subject: [PATCH] fix(sessions): validate HTTP status before parsing Vertex AI session responses --- .../com/google/adk/sessions/ApiResponse.java | 3 ++ .../google/adk/sessions/HttpApiResponse.java | 6 +++ .../adk/sessions/VertexAiApiException.java | 38 +++++++++++++++ .../google/adk/sessions/VertexAiClient.java | 25 +++++++--- .../adk/sessions/VertexAiSessionService.java | 6 +++ .../google/adk/sessions/MockApiAnswer.java | 13 +++++- .../sessions/VertexAiSessionServiceTest.java | 46 ++++++++++++++----- 7 files changed, 117 insertions(+), 20 deletions(-) create mode 100644 core/src/main/java/com/google/adk/sessions/VertexAiApiException.java diff --git a/core/src/main/java/com/google/adk/sessions/ApiResponse.java b/core/src/main/java/com/google/adk/sessions/ApiResponse.java index 7e3393d7d..c43c0d679 100644 --- a/core/src/main/java/com/google/adk/sessions/ApiResponse.java +++ b/core/src/main/java/com/google/adk/sessions/ApiResponse.java @@ -23,6 +23,9 @@ public abstract class ApiResponse implements AutoCloseable { /** Gets the HttpEntity. */ public abstract ResponseBody getResponseBody(); + /** Gets the HTTP status code of the response. */ + public abstract int getStatusCode(); + @Override public abstract void close(); } diff --git a/core/src/main/java/com/google/adk/sessions/HttpApiResponse.java b/core/src/main/java/com/google/adk/sessions/HttpApiResponse.java index f98b7a173..158c4ce7f 100644 --- a/core/src/main/java/com/google/adk/sessions/HttpApiResponse.java +++ b/core/src/main/java/com/google/adk/sessions/HttpApiResponse.java @@ -35,6 +35,12 @@ public ResponseBody getResponseBody() { return response.body(); } + /** Returns the HTTP status code from the response. */ + @Override + public int getStatusCode() { + return response.code(); + } + /** Closes the Http response. */ @Override public void close() { diff --git a/core/src/main/java/com/google/adk/sessions/VertexAiApiException.java b/core/src/main/java/com/google/adk/sessions/VertexAiApiException.java new file mode 100644 index 000000000..6c899f5fb --- /dev/null +++ b/core/src/main/java/com/google/adk/sessions/VertexAiApiException.java @@ -0,0 +1,38 @@ +/* + * Copyright 2025 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.adk.sessions; + +/** + * Signals a non-2xx, non-404 HTTP status from the Vertex AI Session API. Extends {@link + * SessionException} so existing {@code catch (SessionException)} call sites still work. + */ +public final class VertexAiApiException extends SessionException { + private final int statusCode; + + VertexAiApiException(int statusCode, String responseBody) { + super( + "Vertex AI Session API request failed with HTTP status " + + statusCode + + (responseBody == null || responseBody.isEmpty() ? "" : ": " + responseBody)); + this.statusCode = statusCode; + } + + /** Returns the HTTP status code returned by the API. */ + public int statusCode() { + return statusCode; + } +} 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 478b2761b..d6e74fe1c 100644 --- a/core/src/main/java/com/google/adk/sessions/VertexAiClient.java +++ b/core/src/main/java/com/google/adk/sessions/VertexAiClient.java @@ -206,16 +206,29 @@ private Single performApiRequest(String method, String path, String */ @Nullable private static Maybe getJsonResponse(ApiResponse apiResponse) { + if (apiResponse == null) { + return Maybe.empty(); + } try { - if (apiResponse == null || apiResponse.getResponseBody() == null) { + int statusCode = apiResponse.getStatusCode(); + String responseString; + try { + ResponseBody responseBody = apiResponse.getResponseBody(); + responseString = responseBody == null ? "" : responseBody.string(); + } catch (IOException e) { + return Maybe.error(new UncheckedIOException(e)); + } + + if (statusCode == 404) { + return Maybe.empty(); + } + if (statusCode < 200 || statusCode >= 300) { + return Maybe.error(new VertexAiApiException(statusCode, responseString)); + } + if (responseString.isEmpty()) { return Maybe.empty(); } try { - ResponseBody responseBody = apiResponse.getResponseBody(); - String responseString = responseBody.string(); // Read body here - if (responseString.isEmpty()) { - return Maybe.empty(); - } return Maybe.just(objectMapper.readTree(responseString)); } catch (IOException e) { return Maybe.error(new UncheckedIOException(e)); 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..6f61f12f5 100644 --- a/core/src/main/java/com/google/adk/sessions/VertexAiSessionService.java +++ b/core/src/main/java/com/google/adk/sessions/VertexAiSessionService.java @@ -193,6 +193,12 @@ private ListEventsResponse parseListEventsResponse(JsonNode listEventsResponse) .build(); } + /** + * {@inheritDoc} + * + *

On a non-2xx, non-404 HTTP response the returned {@link Maybe} emits a {@link + * VertexAiApiException}. + */ @Override public Maybe getSession( String appName, String userId, String sessionId, Optional config) { 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 36b3d92b9..84c860996 100644 --- a/core/src/test/java/com/google/adk/sessions/MockApiAnswer.java +++ b/core/src/test/java/com/google/adk/sessions/MockApiAnswer.java @@ -99,10 +99,19 @@ public ApiResponse answer(InvocationOnMock invocation) throws Throwable { } private static ApiResponse responseWithBody(String body) { + return responseWithStatus(200, body); + } + + static ApiResponse responseWithStatus(int statusCode, String body) { return new ApiResponse() { @Override public ResponseBody getResponseBody() { - return ResponseBody.create(JSON_MEDIA_TYPE, body); + return body == null ? null : ResponseBody.create(JSON_MEDIA_TYPE, body); + } + + @Override + public int getStatusCode() { + return statusCode; } @Override @@ -144,7 +153,7 @@ private ApiResponse handleGetSession(String path) throws Exception { if (sessionData != null) { return responseWithBody(sessionData); } else { - throw new RuntimeException("Session not found: " + sessionId); + return responseWithStatus(404, ""); } } 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 db3556956..bd703314e 100644 --- a/core/src/test/java/com/google/adk/sessions/VertexAiSessionServiceTest.java +++ b/core/src/test/java/com/google/adk/sessions/VertexAiSessionServiceTest.java @@ -241,15 +241,10 @@ public void createSession_noState_success() throws Exception { } @Test - public void getEmptySession_success() { - RuntimeException exception = - assertThrows( - RuntimeException.class, - () -> - vertexAiSessionService - .getSession("123", "user", "0", Optional.empty()) - .blockingGet()); - assertThat(exception).hasMessageThat().contains("Session not found: 0"); + public void getEmptySession_returnsNull() { + Session session = + vertexAiSessionService.getSession("123", "user", "0", Optional.empty()).blockingGet(); + assertThat(session).isNull(); } @Test @@ -258,14 +253,41 @@ public void getAndDeleteSession_success() throws Exception { vertexAiSessionService.getSession("123", "user", "1", Optional.empty()).blockingGet(); assertThat(session.toJson()).isEqualTo(getMockSession().toJson()); vertexAiSessionService.deleteSession("123", "user", "1").blockingAwait(); - RuntimeException exception = + Session sessionAfterDelete = + vertexAiSessionService.getSession("123", "user", "1", Optional.empty()).blockingGet(); + assertThat(sessionAfterDelete).isNull(); + } + + @Test + public void getSession_permissionDenied_propagatesAsError() { + when(mockApiClient.request(eq("GET"), eq("reasoningEngines/123/sessions/1"), eq(""))) + .thenReturn( + MockApiAnswer.responseWithStatus( + 403, "{\"userId\": \"user\", \"error\": \"permission denied\"}")); + + VertexAiApiException exception = + assertThrows( + VertexAiApiException.class, + () -> + vertexAiSessionService + .getSession("123", "user", "1", Optional.empty()) + .blockingGet()); + assertThat(exception.statusCode()).isEqualTo(403); + } + + @Test + public void getSession_serverError_propagatesAsError() { + when(mockApiClient.request(eq("GET"), eq("reasoningEngines/123/sessions/1"), eq(""))) + .thenReturn(MockApiAnswer.responseWithStatus(500, "{\"error\": \"internal\"}")); + + VertexAiApiException exception = assertThrows( - RuntimeException.class, + VertexAiApiException.class, () -> vertexAiSessionService .getSession("123", "user", "1", Optional.empty()) .blockingGet()); - assertThat(exception).hasMessageThat().contains("Session not found: 1"); + assertThat(exception.statusCode()).isEqualTo(500); } @Test