fix: Handle NOT_FOUND error when fetching Vertex AI sessions

Co-authored-by: Liang Wu <wuliang@google.com>
PiperOrigin-RevId: 855464349

Author: wuliang229

src/google/adk/sessions/vertex_ai_session_service.py

@@ -24,6 +24,7 @@
 from typing import Union
 
 from google.genai import types
+from google.genai.errors import ClientError
 from typing_extensions import override
 
 if TYPE_CHECKING:
@@ -155,13 +156,22 @@ async def get_session(
             )
         }
 
-      get_session_response, events_iterator = await asyncio.gather(
-          api_client.agent_engines.sessions.get(name=session_resource_name),
-          api_client.agent_engines.sessions.events.list(
-              name=session_resource_name,
-              **list_events_kwargs,
-          ),
-      )
+      try:
+        get_session_response, events_iterator = await asyncio.gather(
+            api_client.agent_engines.sessions.get(name=session_resource_name),
+            api_client.agent_engines.sessions.events.list(
+                name=session_resource_name,
+                **list_events_kwargs,
+            ),
+        )
+      except ClientError as e:
+        if e.code == 404:
+          logger.debug(
+              'Session %s not found in Vertex AI Agent Engine.',
+              session_resource_name,
+          )
+          return None
+        raise
       if get_session_response.user_id != user_id:
         raise ValueError(
             f'Session {session_id} does not belong to user {user_id}.'

tests/unittests/sessions/test_vertex_ai_session_service.py

@@ -32,6 +32,7 @@
 from google.adk.sessions.vertex_ai_session_service import VertexAiSessionService
 from google.api_core import exceptions as api_core_exceptions
 from google.genai import types as genai_types
+from google.genai.errors import ClientError
 import pytest
 
 MOCK_SESSION_JSON_1 = {
@@ -552,6 +553,31 @@ async def test_initialize_with_project_location_and_api_key_error():
   )
 
 
+@pytest.mark.asyncio
+@pytest.mark.usefixtures('mock_get_api_client')
+async def test_get_session_returns_none_when_invalid_argument(
+    mock_api_client_instance,
+):
+  session_service = mock_vertex_ai_session_service()
+  # Simulate the API raising a session not found exception.
+  mock_api_client_instance.agent_engines.sessions.get.side_effect = ClientError(
+      code=404,
+      response_json={
+          'message': (
+              'Session (projectNumber: 123, reasoningEngineId: 123, sessionId:'
+              ' 123) not found.'
+          )
+      },
+      response=None,
+  )
+
+  session = await session_service.get_session(
+      app_name='123', user_id='user', session_id='missing'
+  )
+
+  assert session is None
+
+
 @pytest.mark.asyncio
 @pytest.mark.usefixtures('mock_get_api_client')
 @pytest.mark.parametrize('agent_engine_id', [None, '123'])