fix: Use async with for API client in Vertex session management to ensure proper resource management

PiperOrigin-RevId: 830485639

Author: yinghsienwu

src/google/adk/sessions/vertex_ai_session_service.py

@@ -107,19 +107,18 @@ async def create_session(
       )
 
     reasoning_engine_id = self._get_reasoning_engine_id(app_name)
-    api_client = self._get_api_client()
 
     config = {'session_state': state} if state else {}
     config.update(kwargs)
-
-    api_response = await api_client.aio.agent_engines.sessions.create(
-        name=f'reasoningEngines/{reasoning_engine_id}',
-        user_id=user_id,
-        config=config,
-    )
-    logger.debug('Create session response: %s', api_response)
-    get_session_response = api_response.response
-    session_id = get_session_response.name.split('/')[-1]
+    async with self._get_api_client() as api_client:
+      api_response = await api_client.agent_engines.sessions.create(
+          name=f'reasoningEngines/{reasoning_engine_id}',
+          user_id=user_id,
+          config=config,
+      )
+      logger.debug('Create session response: %s', api_response)
+      get_session_response = api_response.response
+      session_id = get_session_response.name.split('/')[-1]
 
     session = Session(
         app_name=app_name,
@@ -140,30 +139,29 @@ async def get_session(
       config: Optional[GetSessionConfig] = None,
   ) -> Optional[Session]:
     reasoning_engine_id = self._get_reasoning_engine_id(app_name)
-    api_client = self._get_api_client()
     session_resource_name = (
         f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}'
     )
-
-    # Get session resource and events in parallel.
-    list_events_kwargs = {}
-    if config and not config.num_recent_events and config.after_timestamp:
-      # Filter events based on timestamp.
-      list_events_kwargs['config'] = {
-          'filter': 'timestamp>="{}"'.format(
-              datetime.datetime.fromtimestamp(
-                  config.after_timestamp, tz=datetime.timezone.utc
-              ).isoformat()
-          )
-      }
-
-    get_session_response, events_iterator = await asyncio.gather(
-        api_client.aio.agent_engines.sessions.get(name=session_resource_name),
-        api_client.aio.agent_engines.sessions.events.list(
-            name=session_resource_name,
-            **list_events_kwargs,
-        ),
-    )
+    async with self._get_api_client() as api_client:
+      # Get session resource and events in parallel.
+      list_events_kwargs = {}
+      if config and not config.num_recent_events and config.after_timestamp:
+        # Filter events based on timestamp.
+        list_events_kwargs['config'] = {
+            'filter': 'timestamp>="{}"'.format(
+                datetime.datetime.fromtimestamp(
+                    config.after_timestamp, tz=datetime.timezone.utc
+                ).isoformat()
+            )
+        }
+
+      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,
+          ),
+      )
 
     if get_session_response.user_id != user_id:
       raise ValueError(
@@ -196,51 +194,52 @@ async def list_sessions(
       self, *, app_name: str, user_id: Optional[str] = None
   ) -> ListSessionsResponse:
     reasoning_engine_id = self._get_reasoning_engine_id(app_name)
-    api_client = self._get_api_client()
 
-    sessions = []
-    config = {}
-    if user_id is not None:
-      config['filter'] = f'user_id="{user_id}"'
-    sessions_iterator = await api_client.aio.agent_engines.sessions.list(
-        name=f'reasoningEngines/{reasoning_engine_id}',
-        config=config,
-    )
-
-    for api_session in sessions_iterator:
-      sessions.append(
-          Session(
-              app_name=app_name,
-              user_id=api_session.user_id,
-              id=api_session.name.split('/')[-1],
-              state=getattr(api_session, 'session_state', None) or {},
-              last_update_time=api_session.update_time.timestamp(),
-          )
+    async with self._get_api_client() as api_client:
+      sessions = []
+      config = {}
+      if user_id is not None:
+        config['filter'] = f'user_id="{user_id}"'
+      sessions_iterator = await api_client.agent_engines.sessions.list(
+          name=f'reasoningEngines/{reasoning_engine_id}',
+          config=config,
       )
 
+      for api_session in sessions_iterator:
+        sessions.append(
+            Session(
+                app_name=app_name,
+                user_id=api_session.user_id,
+                id=api_session.name.split('/')[-1],
+                state=getattr(api_session, 'session_state', None) or {},
+                last_update_time=api_session.update_time.timestamp(),
+            )
+        )
+
     return ListSessionsResponse(sessions=sessions)
 
   async def delete_session(
       self, *, app_name: str, user_id: str, session_id: str
   ) -> None:
     reasoning_engine_id = self._get_reasoning_engine_id(app_name)
-    api_client = self._get_api_client()
 
-    try:
-      await api_client.aio.agent_engines.sessions.delete(
-          name=f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}',
-      )
-    except Exception as e:
-      logger.error('Error deleting session %s: %s', session_id, e)
-      raise e
+    async with self._get_api_client() as api_client:
+      try:
+        await api_client.agent_engines.sessions.delete(
+            name=(
+                f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}'
+            ),
+        )
+      except Exception as e:
+        logger.error('Error deleting session %s: %s', session_id, e)
+        raise
 
   @override
   async def append_event(self, session: Session, event: Event) -> Event:
     # Update the in-memory session.
     await super().append_event(session=session, event=event)
 
     reasoning_engine_id = self._get_reasoning_engine_id(session.app_name)
-    api_client = self._get_api_client()
 
     config = {}
     if event.content:
@@ -284,15 +283,16 @@ async def append_event(self, session: Session, event: Event) -> Event:
       )
     config['event_metadata'] = metadata_dict
 
-    await api_client.aio.agent_engines.sessions.events.append(
-        name=f'reasoningEngines/{reasoning_engine_id}/sessions/{session.id}',
-        author=event.author,
-        invocation_id=event.invocation_id,
-        timestamp=datetime.datetime.fromtimestamp(
-            event.timestamp, tz=datetime.timezone.utc
-        ),
-        config=config,
-    )
+    async with self._get_api_client() as api_client:
+      await api_client.agent_engines.sessions.events.append(
+          name=f'reasoningEngines/{reasoning_engine_id}/sessions/{session.id}',
+          author=event.author,
+          invocation_id=event.invocation_id,
+          timestamp=datetime.datetime.fromtimestamp(
+              event.timestamp, tz=datetime.timezone.utc
+          ),
+          config=config,
+      )
     return event
 
   def _get_reasoning_engine_id(self, app_name: str):
@@ -318,7 +318,7 @@ def _api_client_http_options_override(
   ) -> Optional[Union[types.HttpOptions, types.HttpOptionsDict]]:
     return None
 
-  def _get_api_client(self) -> vertexai.Client:
+  def _get_api_client(self) -> vertexai.AsyncClient:
     """Instantiates an API client for the given project and location.
 
     Returns:
@@ -331,7 +331,7 @@ def _get_api_client(self) -> vertexai.Client:
         location=self._location,
         http_options=self._api_client_http_options_override(),
         api_key=self._express_mode_api_key,
-    )
+    ).aio
 
 
 def _from_api_event(api_event_obj: vertexai.types.SessionEvent) -> Event:

tests/unittests/sessions/test_vertex_ai_session_service.py

@@ -228,24 +228,30 @@ def _convert_to_object(data):
     return data
 
 
-class MockApiClient:
+class MockAsyncClient:
   """Mocks the API Client."""
 
   def __init__(self) -> None:
     """Initializes MockClient."""
     self.session_dict: dict[str, Any] = {}
     self.event_dict: dict[str, Tuple[List[Any], Optional[str]]] = {}
-    self.aio = mock.Mock()
-    self.aio.agent_engines.sessions.get.side_effect = self._get_session
-    self.aio.agent_engines.sessions.list.side_effect = self._list_sessions
-    self.aio.agent_engines.sessions.delete.side_effect = self._delete_session
-    self.aio.agent_engines.sessions.create.side_effect = self._create_session
-    self.aio.agent_engines.sessions.events.list.side_effect = self._list_events
-    self.aio.agent_engines.sessions.events.append.side_effect = (
-        self._append_event
-    )
+    self.agent_engines = mock.AsyncMock()
+    self.agent_engines.sessions.get.side_effect = self._get_session
+    self.agent_engines.sessions.list.side_effect = self._list_sessions
+    self.agent_engines.sessions.delete.side_effect = self._delete_session
+    self.agent_engines.sessions.create.side_effect = self._create_session
+    self.agent_engines.sessions.events.list.side_effect = self._list_events
+    self.agent_engines.sessions.events.append.side_effect = self._append_event
     self.last_create_session_config: dict[str, Any] = {}
 
+  async def __aenter__(self):
+    """Enters the asynchronous context."""
+    return self
+
+  async def __aexit__(self, exc_type, exc_val, exc_tb):
+    """Exits the asynchronous context."""
+    pass
+
   async def _get_session(self, name: str):
     session_id = name.split('/')[-1]
     if session_id in self.session_dict:
@@ -374,7 +380,7 @@ def mock_vertex_ai_session_service(
 @pytest.fixture
 def mock_api_client_instance():
   """Creates a mock API client instance for testing."""
-  api_client = MockApiClient()
+  api_client = MockAsyncClient()
   api_client.session_dict = {
       '1': MOCK_SESSION_JSON_1,
       '2': MOCK_SESSION_JSON_2,