feat: Populate X-Server-Timeout header when a request timeout is set.

PiperOrigin-RevId: 743955769

Author: jscud

google/genai/_api_client.py

@@ -27,6 +27,7 @@
 import io
 import json
 import logging
+import math
 import os
 import sys
 import time
@@ -97,6 +98,14 @@ def _patch_http_options(
   return copy_option
 
 
+def _populate_server_timeout_header(
+    headers: dict[str, str], timeout_in_seconds: Optional[Union[float, int]]
+) -> None:
+  """Populates the server timeout header in the headers dict."""
+  if timeout_in_seconds and 'X-Server-Timeout' not in headers:
+    headers['X-Server-Timeout'] = str(math.ceil(timeout_in_seconds))
+
+
 def _join_url_path(base_url: str, path: str) -> str:
   parsed_base = urlparse(base_url)
   base_path = (
@@ -539,6 +548,9 @@ def _build_request(
 
     if patched_http_options.headers is None:
       raise ValueError('Request headers must be set.')
+    _populate_server_timeout_header(
+        patched_http_options.headers, timeout_in_seconds
+    )
     return HttpRequest(
         method=http_method,
         url=url,
@@ -795,14 +807,16 @@ def _upload_fd(
             else self._http_options.timeout
         )
       timeout_in_seconds = _get_timeout_in_seconds(timeout)
+      upload_headers = {
+          'X-Goog-Upload-Command': upload_command,
+          'X-Goog-Upload-Offset': str(offset),
+          'Content-Length': str(chunk_size),
+      }
+      _populate_server_timeout_header(upload_headers, timeout_in_seconds)
       response = self._httpx_client.request(
           method='POST',
           url=upload_url,
-          headers={
-              'X-Goog-Upload-Command': upload_command,
-              'X-Goog-Upload-Offset': str(offset),
-              'Content-Length': str(chunk_size),
-          },
+          headers=upload_headers,
           content=file_chunk,
           timeout=timeout_in_seconds,
       )
@@ -941,15 +955,17 @@ async def _async_upload_fd(
             else self._http_options.timeout
         )
       timeout_in_seconds = _get_timeout_in_seconds(timeout)
+      upload_headers = {
+          'X-Goog-Upload-Command': upload_command,
+          'X-Goog-Upload-Offset': str(offset),
+          'Content-Length': str(chunk_size),
+      }
+      _populate_server_timeout_header(upload_headers, timeout_in_seconds)
       response = await self._async_httpx_client.request(
           method='POST',
           url=upload_url,
           content=file_chunk,
-          headers={
-              'X-Goog-Upload-Command': upload_command,
-              'X-Goog-Upload-Offset': str(offset),
-              'Content-Length': str(chunk_size),
-          },
+          headers=upload_headers,
           timeout=timeout_in_seconds,
       )
       offset += chunk_size

google/genai/tests/client/test_http_options.py

@@ -77,3 +77,63 @@ def test_patch_http_options_appends_version_headers():
   patched = _api_client._patch_http_options(original_options, patch_options)
   assert 'user-agent' in patched.headers
   assert 'x-goog-api-client' in patched.headers
+
+
+def test_setting_timeout_populates_server_timeout_header():
+  api_client = _api_client.BaseApiClient(
+      vertexai=False,
+      api_key='test_api_key',
+      http_options=types.HttpOptions(timeout=10000),
+  )
+  request = api_client._build_request(
+      http_method='POST',
+      path='sample/path',
+      request_dict={},
+  )
+  assert 'X-Server-Timeout' in request.headers
+  assert request.headers['X-Server-Timeout'] == '10'
+
+
+def test_timeout_rounded_to_nearest_second():
+  api_client = _api_client.BaseApiClient(
+      vertexai=False,
+      api_key='test_api_key',
+  )
+  http_options = types.HttpOptions(timeout=7300)
+  request = api_client._build_request(
+      http_method='POST',
+      path='sample/path',
+      request_dict={},
+      http_options=http_options,
+  )
+  assert request.headers['X-Server-Timeout'] == '8'
+
+
+def test_server_timeout_not_overwritten():
+  api_client = _api_client.BaseApiClient(
+      vertexai=False,
+      api_key='test_api_key',
+  )
+  http_options = types.HttpOptions(
+      headers={'X-Server-Timeout': '3'},
+      timeout=11000)
+  request = api_client._build_request(
+      http_method='POST',
+      path='sample/path',
+      request_dict={},
+      http_options=http_options,
+  )
+  assert request.headers['X-Server-Timeout'] == '3'
+
+
+def test_server_timeout_not_set_by_default():
+  api_client = _api_client.BaseApiClient(
+      vertexai=False,
+      api_key='test_api_key',
+  )
+  request = api_client._build_request(
+      http_method='POST',
+      path='sample/path',
+      request_dict={},
+  )
+  assert not 'X-Server-Timeout' in request.headers