fix: Pass service_tier to Agent Platform backend

google/genai/_extra_utils.py

@@ -777,3 +777,27 @@ def get_usage_header(
 
   http_options.headers = existing_headers
   return config_model
+
+
+def set_agent_platform_service_tier_header(
+    config: Optional[types.GenerateContentConfig],
+    is_vertex: bool,
+) -> None:
+  """Injects the service_tier as an HTTP header for Gemini Enterprise Agent Platform requests."""
+  if not is_vertex or not config or not config.service_tier:
+    return
+
+  # Extract the raw string value (e.g. 'flex', 'standard', 'priority')
+  tier_val = (
+      config.service_tier.value
+      if hasattr(config.service_tier, 'value')
+      else str(config.service_tier)
+  )
+  tier_val = tier_val.lower().replace('service_tier_', '')
+
+  if config.http_options is None:
+    config.http_options = types.HttpOptions(headers={})
+  if config.http_options.headers is None:
+    config.http_options.headers = {}
+
+  config.http_options.headers['X-Vertex-AI-LLM-Shared-Request-Type'] = tier_val

google/genai/models.py

@@ -1615,9 +1615,6 @@ def _GenerateContentConfig_to_vertex(
         getv(from_object, ['model_armor_config']),
     )
 
-  if getv(from_object, ['service_tier']) is not None:
-    setv(parent_object, ['serviceTier'], getv(from_object, ['service_tier']))
-
   if getv(from_object, ['audio_transcription_config']) is not None:
     setv(
         to_object,
@@ -6497,6 +6494,11 @@ def generate_content(
         _extra_utils.find_afc_incompatible_tool_indexes(config)
     )
     parsed_config = _extra_utils.parse_config_for_mcp_usage(config)
+
+    _extra_utils.set_agent_platform_service_tier_header(
+        parsed_config, is_vertex=getattr(self._api_client, 'vertexai', False)
+    )
+
     if (
         parsed_config
         and parsed_config.tools
@@ -6663,6 +6665,11 @@ def generate_content_stream(
         _extra_utils.find_afc_incompatible_tool_indexes(config)
     )
     parsed_config = _extra_utils.parse_config_for_mcp_usage(config)
+
+    _extra_utils.set_agent_platform_service_tier_header(
+        parsed_config, is_vertex=getattr(self._api_client, 'vertexai', False)
+    )
+
     if (
         parsed_config
         and parsed_config.tools
@@ -8628,6 +8635,10 @@ async def generate_content(
     else:
       parsed_config = config.model_copy(deep=True)
 
+    _extra_utils.set_agent_platform_service_tier_header(
+        parsed_config, is_vertex=getattr(self._api_client, 'vertexai', False)
+    )
+
     # Use AsyncExitStack to keep MCP connections alive across the entire AFC loop
     async with contextlib.AsyncExitStack() as stack:
 
@@ -8859,6 +8870,10 @@ async def generate_content_stream(
     else:
       parsed_config = config.model_copy(deep=True)
 
+    _extra_utils.set_agent_platform_service_tier_header(
+        parsed_config, is_vertex=getattr(self._api_client, 'vertexai', False)
+    )
+
     incompatible_tools_indexes = (
         _extra_utils.find_afc_incompatible_tool_indexes(
             parsed_config,

google/genai/tests/models/test_generate_content.py

@@ -577,24 +577,22 @@ class InstrumentEnum(Enum):
     pytest_helper.TestTableItem(
         name='test_service_tier',
         parameters=types._GenerateContentParameters(
-            model=GEMINI_FLASH_LATEST,
+            model='gemini-3.1-pro-preview',
             contents=t.t_contents('What is your name?'),
             config={
                 'service_tier': 'FLEX',
             },
         ),
-        exception_if_vertex='400',
     ),
     pytest_helper.TestTableItem(
         name='test_service_tier_lower',
         parameters=types._GenerateContentParameters(
-            model=GEMINI_FLASH_LATEST,
+            model='gemini-3.1-pro-preview',
             contents=t.t_contents('What is your name?'),
             config={
                 'service_tier': 'flex',
             },
         ),
-        exception_if_vertex='400',
     ),
     pytest_helper.TestTableItem(
         name='test_audio_transcription_config',