From 207a1dd74ed5d9590c7b260d34ccde00087dd5ac Mon Sep 17 00:00:00 2001 From: Ayush Agrawal Date: Wed, 21 Oct 2026 00:28:48 -0700 Subject: [PATCH] fix: Stop per-request Authorization and X-Server-Timeout from leaking into client default headers fixes #3158 PiperOrigin-RevId: 1003716874 --- google/genai/_api_client.py | 8 +- google/genai/live.py | 2 + .../genai/tests/client/test_http_options.py | 72 +++++++++++++++ google/genai/tests/interactions/test_auth.py | 92 +++++++++++++++++++ google/genai/tests/live/test_live.py | 73 +++++++++++++++ 5 files changed, 243 insertions(+), 4 deletions(-) diff --git a/google/genai/_api_client.py b/google/genai/_api_client.py index 82e48973c..8a0618c7d 100644 --- a/google/genai/_api_client.py +++ b/google/genai/_api_client.py @@ -1547,13 +1547,13 @@ 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 - ) + # Copy so per-request headers do not become client defaults. + headers = dict(patched_http_options.headers) + populate_server_timeout_header(headers, timeout_in_seconds) return HttpRequest( method=http_method, url=url, - headers=patched_http_options.headers, + headers=headers, data=request_dict, timeout=timeout_in_seconds, ) diff --git a/google/genai/live.py b/google/genai/live.py index 78b173a33..2905d8946 100644 --- a/google/genai/live.py +++ b/google/genai/live.py @@ -1142,6 +1142,8 @@ async def connect( ) if not headers.get('Authorization'): headers['Authorization'] = f'Bearer {bearer_token}' + if creds.quota_project_id and not headers.get('x-goog-user-project'): + headers['x-goog-user-project'] = creds.quota_project_id location = self._api_client.location project = self._api_client.project diff --git a/google/genai/tests/client/test_http_options.py b/google/genai/tests/client/test_http_options.py index f15cde97a..c5c08de54 100644 --- a/google/genai/tests/client/test_http_options.py +++ b/google/genai/tests/client/test_http_options.py @@ -16,6 +16,8 @@ """Tests for client behavior when issuing requests.""" +from unittest import mock + import httpx import pytest @@ -211,3 +213,73 @@ def test_setting_30m_timeout_populates_header_and_request_timeout(): ) assert request.headers['X-Server-Timeout'] == '1800' assert request.timeout == 1800.0 + + +def test_build_request_does_not_mutate_client_headers(): + # Per-request headers must not leak into the client's default headers. + api_client = _api_client.BaseApiClient( + vertexai=False, + api_key='test_api_key', + http_options=types.HttpOptions(timeout=10000), + ) + headers_before = dict(api_client._http_options.headers) + request = api_client._build_request( + http_method='POST', + path='sample/path', + request_dict={}, + ) + assert request.headers['X-Server-Timeout'] == '10' + assert request.headers is not api_client._http_options.headers + assert api_client._http_options.headers == headers_before + + +def test_per_request_timeout_not_overridden_by_earlier_request(): + # A client-timeout request must not pin X-Server-Timeout for later ones. + api_client = _api_client.BaseApiClient( + vertexai=False, + api_key='test_api_key', + http_options=types.HttpOptions(timeout=10000), + ) + api_client._build_request( + http_method='POST', + path='sample/path', + request_dict={}, + ) + request = api_client._build_request( + http_method='POST', + path='sample/path', + request_dict={}, + http_options=types.HttpOptions(timeout=3000), + ) + assert request.headers['X-Server-Timeout'] == '3' + + +def test_request_does_not_leak_authorization_into_client_headers(): + # The bearer token and quota project are set per request only. + creds = mock.Mock() + creds.quota_project_id = 'test-quota-project' + api_client = _api_client.BaseApiClient( + vertexai=True, + project='test-project', + location='us-central1', + credentials=creds, + ) + with ( + mock.patch.object( + _api_client.BaseApiClient, '_access_token', return_value='token-1' + ), + mock.patch.object( + httpx.Client, + 'send', + return_value=httpx.Response( + 200, request=httpx.Request('POST', ''), json={} + ), + ) as mock_send, + ): + api_client.request('post', 'sample/path', {}) + + sent_headers = mock_send.call_args[0][0].headers + assert sent_headers['authorization'] == 'Bearer token-1' + assert sent_headers['x-goog-user-project'] == 'test-quota-project' + assert 'Authorization' not in api_client._http_options.headers + assert 'x-goog-user-project' not in api_client._http_options.headers diff --git a/google/genai/tests/interactions/test_auth.py b/google/genai/tests/interactions/test_auth.py index c96e7592f..3873dbde5 100644 --- a/google/genai/tests/interactions/test_auth.py +++ b/google/genai/tests/interactions/test_auth.py @@ -484,3 +484,95 @@ async def test_async_interactions_vertex_extra_headers_override(): headers = mock_send.call_args[0][0].headers assert headers['x-goog-api-key'] == 'manual-key' assert 'authorization' not in headers + + +_GENERATE_CONTENT_BODY = { + 'candidates': [{ + 'content': {'role': 'model', 'parts': [{'text': 'ok'}]}, + 'finishReason': 'STOP', + }] +} + + +def _send_recording_auth(recorded_auth): + """Returns a fake send that records auth and answers REST or Interactions.""" + + def send(req, *args, **kwargs): + recorded_auth.append((req.url.path, req.headers.get('authorization'))) + if '/interactions' in req.url.path: + return Response( + 200, + request=Request('POST', ''), + headers={'content-type': 'application/json'}, + content='{"status": "completed"}', + ) + return Response( + 200, request=Request('POST', ''), json=_GENERATE_CONTENT_BODY + ) + + return send + + +def test_interactions_vertex_fresh_token_after_models_call(): + # A REST call before the first Interactions access must not pin its token. + creds = mock.Mock() + creds.quota_project_id = 'test-quota-project' + client = Client( + vertexai=True, + project='test-project', + location='us-central1', + credentials=creds, + ) + recorded_auth = [] + + with ( + mock.patch.object( + BaseApiClient, '_access_token', side_effect=['token-1', 'token-2'] + ), + mock.patch.object( + HTTPClient, 'send', side_effect=_send_recording_auth(recorded_auth) + ), + ): + client.models.generate_content(model='gemini-2.5-flash', contents='hi') + client.interactions.create(model='gemini-2.5-flash', input='Hello') + + assert len(recorded_auth) == 2 + assert '/interactions' not in recorded_auth[0][0] + assert recorded_auth[0][1] == 'Bearer token-1' + assert recorded_auth[1][0].endswith('/interactions') + assert recorded_auth[1][1] == 'Bearer token-2' + + +@pytest.mark.asyncio +async def test_async_interactions_vertex_fresh_token_after_models_call(): + # Async twin; client.aio shares the sync client's BaseApiClient. + creds = mock.Mock() + creds.quota_project_id = 'test-quota-project' + client = Client( + vertexai=True, + project='test-project', + location='us-central1', + credentials=creds, + ) + recorded_auth = [] + sync_send = _send_recording_auth(recorded_auth) + + async def async_send(req, *args, **kwargs): + return sync_send(req, *args, **kwargs) + + with ( + mock.patch.object( + BaseApiClient, '_access_token', side_effect=['token-1', 'token-2'] + ), + mock.patch.object(HTTPClient, 'send', side_effect=sync_send), + mock.patch.object(AsyncHttpxClient, 'send', side_effect=async_send), + ): + client.models.generate_content(model='gemini-2.5-flash', contents='hi') + await client.aio.interactions.create( + model='gemini-2.5-flash', input='Hello' + ) + + assert len(recorded_auth) == 2 + assert recorded_auth[0][1] == 'Bearer token-1' + assert recorded_auth[1][0].endswith('/interactions') + assert recorded_auth[1][1] == 'Bearer token-2' diff --git a/google/genai/tests/live/test_live.py b/google/genai/tests/live/test_live.py index 8062b868d..692aaa68e 100644 --- a/google/genai/tests/live/test_live.py +++ b/google/genai/tests/live/test_live.py @@ -31,6 +31,7 @@ import certifi from google.oauth2.credentials import Credentials +import httpx import pytest from websockets import client @@ -2214,6 +2215,78 @@ async def _test_connect(): await _test_connect() +@pytest.mark.asyncio +async def test_connect_uses_current_token_after_rest_call(mock_websocket): + # A REST call must not pin its bearer token for later Live connections. + credentials = Credentials(token='token-1') + client = gl_client.BaseApiClient( + vertexai=True, + project='test_project', + location='us-central1', + credentials=credentials, + ) + capture = {} + + @contextlib.asynccontextmanager + async def mock_connect(uri, additional_headers=None, **kwargs): + capture['headers'] = additional_headers + yield mock_websocket + + with patch.object( + httpx.Client, + 'send', + return_value=httpx.Response( + 200, request=httpx.Request('POST', ''), json={} + ), + ): + client.request('post', 'sample/path', {}) + credentials.token = 'token-2' + + with patch.object(live, 'ws_connect', new=mock_connect): + async with live.AsyncLive(client).connect(model='test-model'): + pass + + assert capture['headers']['Authorization'] == 'Bearer token-2' + + +@pytest.mark.asyncio +async def test_connect_sets_user_project_from_credentials(mock_websocket): + # Live sends the credentials' quota project, as REST requests do. + credentials = Credentials(token='t', quota_project_id='test-quota-project') + client = mock_api_client(vertexai=True, credentials=credentials) + capture = {} + + @contextlib.asynccontextmanager + async def mock_connect(uri, additional_headers=None, **kwargs): + capture['headers'] = additional_headers + yield mock_websocket + + with patch.object(live, 'ws_connect', new=mock_connect): + async with live.AsyncLive(client).connect(model='test-model'): + pass + + assert capture['headers']['x-goog-user-project'] == 'test-quota-project' + + +@pytest.mark.asyncio +async def test_connect_no_user_project_without_quota_project(mock_websocket): + # Without a quota project on the credentials, Live omits the header. + credentials = Credentials(token='t') + client = mock_api_client(vertexai=True, credentials=credentials) + capture = {} + + @contextlib.asynccontextmanager + async def mock_connect(uri, additional_headers=None, **kwargs): + capture['headers'] = additional_headers + yield mock_websocket + + with patch.object(live, 'ws_connect', new=mock_connect): + async with live.AsyncLive(client).connect(model='test-model'): + pass + + assert 'x-goog-user-project' not in capture['headers'] + + @pytest.mark.asyncio async def test_connect_with_custom_base_url(mock_websocket): # mock api client