chore: Adding adk header to track MB and session service usage via ADK SDK

PiperOrigin-RevId: 938299570

Author: google-genai-bot

src/google/adk/memory/vertex_ai_memory_bank_service.py

@@ -26,6 +26,7 @@
 from google.genai import types
 from typing_extensions import override
 
+from ..utils._google_client_headers import get_tracking_headers
 from ..utils.vertex_ai_utils import get_express_mode_api_key
 from .base_memory_service import BaseMemoryService
 from .base_memory_service import SearchMemoryResponse
@@ -616,9 +617,17 @@ def _get_api_client(self) -> vertexai.AsyncClient:
     """
     import vertexai
 
+    http_options = types.HttpOptions(headers=get_tracking_headers())
     if self._express_mode_api_key:
-      return vertexai.Client(api_key=self._express_mode_api_key).aio
-    return vertexai.Client(project=self._project, location=self._location).aio
+      return vertexai.Client(
+          http_options=http_options,
+          api_key=self._express_mode_api_key,
+      ).aio
+    return vertexai.Client(
+        project=self._project,
+        location=self._location,
+        http_options=http_options,
+    ).aio
 
 
 def _log_ingest_task_error(task: asyncio.Task) -> None:

src/google/adk/sessions/vertex_ai_session_service.py

@@ -22,7 +22,6 @@
 from typing import Any
 from typing import Optional
 from typing import TYPE_CHECKING
-from typing import Union
 
 from google.genai import types
 from google.genai.errors import ClientError
@@ -36,6 +35,7 @@
 from ..events.event import Event
 from ..events.event_actions import EventActions
 from ..events.event_actions import EventCompaction
+from ..utils._google_client_headers import get_tracking_headers
 from ..utils.vertex_ai_utils import get_express_mode_api_key
 from .base_session_service import BaseSessionService
 from .base_session_service import GetSessionConfig
@@ -506,11 +506,6 @@ def _get_reasoning_engine_id(self, app_name: str):
 
     return match.groups()[-1]
 
-  def _api_client_http_options_override(
-      self,
-  ) -> Optional[Union[types.HttpOptions, types.HttpOptionsDict]]:
-    return None
-
   def _get_api_client(self) -> vertexai.AsyncClient:
     """Instantiates an API client for the given project and location.
 
@@ -519,15 +514,16 @@ def _get_api_client(self) -> vertexai.AsyncClient:
     """
     import vertexai
 
+    http_options = types.HttpOptions(headers=get_tracking_headers())
     if self._express_mode_api_key:
       return vertexai.Client(
-          http_options=self._api_client_http_options_override(),
+          http_options=http_options,
           api_key=self._express_mode_api_key,
       ).aio
     return vertexai.Client(
         project=self._project,
         location=self._location,
-        http_options=self._api_client_http_options_override(),
+        http_options=http_options,
     ).aio
 
 

tests/unittests/memory/test_vertex_ai_memory_bank_service.py

@@ -1205,3 +1205,15 @@ async def failing_async_iterator():
 
   assert len(result.memories) == 1
   assert result.memories[0].content.parts[0].text == 'good fact'
+
+
+def test_get_api_client_attaches_tracking_headers():
+  memory_service = mock_vertex_ai_memory_bank_service()
+  with mock.patch('vertexai.Client') as mock_client:
+    _ = memory_service._get_api_client()
+    mock_client.assert_called_once()
+    _, kwargs = mock_client.call_args
+    assert 'http_options' in kwargs
+    http_options = kwargs['http_options']
+    assert 'x-goog-api-client' in http_options.headers
+    assert 'google-adk/' in http_options.headers['x-goog-api-client']

tests/unittests/sessions/test_vertex_ai_session_service.py

@@ -1522,3 +1522,15 @@ async def test_get_session_strips_full_resource_name(
   mock_api_client_instance.agent_engines.sessions.get.assert_called_once_with(
       name='reasoningEngines/123/sessions/session-123'
   )
+
+
+def test_get_api_client_attaches_tracking_headers():
+  session_service = mock_vertex_ai_session_service()
+  with mock.patch('vertexai.Client') as mock_client:
+    _ = session_service._get_api_client()
+    mock_client.assert_called_once()
+    _, kwargs = mock_client.call_args
+    assert 'http_options' in kwargs
+    http_options = kwargs['http_options']
+    assert 'x-goog-api-client' in http_options.headers
+    assert 'google-adk/' in http_options.headers['x-goog-api-client']