Skip to content

Commit 72d1866

Browse files
feat: Update for initializing session in init
feat: Update for initializing session in init
1 parent 38cd9d6 commit 72d1866

1 file changed

Lines changed: 25 additions & 5 deletions

File tree

src/google/adk/integrations/agent_registry/agent_registry.py

Lines changed: 25 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -212,6 +212,22 @@ def __init__(
212212
raise RuntimeError(
213213
f"Failed to get default Google Cloud credentials: {e}"
214214
) from e
215+
# Instantiate and configure AuthorizedSession once during initialization
216+
self._session = requests_auth.AuthorizedSession(
217+
credentials=self._credentials
218+
)
219+
220+
use_client_cert = _use_client_cert_effective()
221+
client_cert_source = None
222+
if use_client_cert:
223+
client_cert_source = (
224+
mtls.default_client_cert_source()
225+
if mtls.has_default_client_cert_source()
226+
else None
227+
)
228+
self._session.configure_mtls_channel(client_cert_source)
229+
230+
self._base_url = _get_agent_registry_base_url(client_cert_source)
215231

216232
def _get_auth_headers(self) -> Dict[str, str]:
217233
"""Refreshes credentials and returns authorization headers."""
@@ -255,15 +271,19 @@ def _make_request(
255271
base_url = _get_agent_registry_base_url(client_cert_source)
256272

257273
if path.startswith("projects/"):
258-
url = f"{base_url}/{path}"
274+
url = f"{self._base_url}/{path}"
259275
else:
260-
url = f"{base_url}/{self._base_path}/{path}"
276+
url = f"{self._base_url}/{self._base_path}/{path}"
277+
headers = {}
278+
quota_project_id = (
279+
getattr(self._credentials, "quota_project_id", None) or self.project_id
280+
)
281+
if quota_project_id:
282+
headers["x-goog-user-project"] = quota_project_id
261283

262284
try:
263285
# Using AuthorizedSession for internal API calls to handle mTLS/Auth.
264-
response = session.get(
265-
url, headers=self._get_auth_headers(), params=params
266-
)
286+
response = self._session.get(url, headers=headers, params=params)
267287
response.raise_for_status()
268288
return response.json()
269289
except requests.exceptions.HTTPError as e:

0 commit comments

Comments
 (0)