fix: keep live music API keys out of websocket urls

Author: he-yufeng

google/genai/live_music.py

@@ -175,10 +175,10 @@ async def connect(self, *, model: str) -> AsyncIterator[AsyncMusicSession]:
     transformed_model = t.t_model(self._api_client, model)
 
     if self._api_client.api_key:
-      api_key = self._api_client.api_key
       version = self._api_client._http_options.api_version
-      uri = f'{base_url}/ws/google.ai.generativelanguage.{version}.GenerativeService.BidiGenerateMusic?key={api_key}'
-      headers = self._api_client._http_options.headers
+      uri = f'{base_url}/ws/google.ai.generativelanguage.{version}.GenerativeService.BidiGenerateMusic'
+      original_headers = self._api_client._http_options.headers
+      headers = original_headers.copy() if original_headers is not None else {}
 
       # Only mldev supported
       request_dict = _common.convert_to_dict(

google/genai/tests/live/test_live_music.py

@@ -133,6 +133,33 @@ async def _test_connect():
   return await _test_connect()
 
 
+@pytest.mark.asyncio
+async def test_connect_uses_header_auth_without_query_key(mock_websocket):
+  api_client = mock_api_client(vertexai=False)
+  api_client._websocket_base_url = lambda: 'wss://generativelanguage.googleapis.com'
+  api_client._http_options.api_version = 'v1beta'
+  api_client._http_options.headers['x-goog-api-key'] = 'TEST_API_KEY'
+  captured = {}
+
+  @contextlib.asynccontextmanager
+  async def mock_connect(uri, additional_headers=None):
+    captured['uri'] = uri
+    captured['headers'] = additional_headers
+    yield mock_websocket
+
+  @patch.object(live_music, 'connect', new=mock_connect)
+  async def _test_connect():
+    live_module = live.AsyncLive(api_client)
+    async with live_module.music.connect(model='test_model'):
+      pass
+
+  await _test_connect()
+
+  assert 'TEST_API_KEY' not in captured['uri']
+  assert '?key=' not in captured['uri']
+  assert captured['headers']['x-goog-api-key'] == 'TEST_API_KEY'
+
+
 def test_mldev_from_env(monkeypatch):
   api_key = 'google_api_key'
   monkeypatch.setenv('GOOGLE_API_KEY', api_key)