Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions packages/toolbox-adk/requirements.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
-e ../toolbox-core
google-auth==2.43.0
google-auth-oauthlib==1.2.1
google-adk==1.20.0
typing-extensions==4.12.2
7 changes: 6 additions & 1 deletion packages/toolbox-adk/src/toolbox_adk/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.

import logging
from contextvars import ContextVar
from typing import Any, Awaitable, Callable, Dict, Optional, Union

Expand Down Expand Up @@ -118,8 +119,12 @@ def get_token() -> str:
try:
token = id_token.fetch_id_token(request, audience)
return f"Bearer {token}"
except Exception:
except Exception as e:
# Fallback to default credentials
logging.warning(
f"Failed to fetch ID token for audience {audience} using ADC: {e}. "
"Falling back to google.auth.default()."
)
creds, _ = google.auth.default()
if not creds.valid:
creds.refresh(request)
Expand Down
7 changes: 7 additions & 0 deletions packages/toolbox-adk/src/toolbox_adk/tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.

import logging
from typing import Any, Awaitable, Callable, Dict, Optional, cast

import toolbox_core
Expand Down Expand Up @@ -146,6 +147,12 @@ async def run_async(
ctx.error = e
if "credential" in str(e).lower() or isinstance(e, ValueError):
raise e

logging.warning(
f"Unexpected error in get_auth_response during 3LO retrieval: {e}. "
"Falling back to request_credential.",
exc_info=True
)
# Fallback to request logic
ctx_any = cast(Any, tool_context)
ctx_any.request_credential(auth_config_adk)
Expand Down
169 changes: 169 additions & 0 deletions packages/toolbox-adk/tests/integration/conftest.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,169 @@
# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Contains pytest fixtures that are accessible from all
files present in the same directory."""

from __future__ import annotations

import os
import platform
import subprocess
import tempfile
import time
from typing import Generator

import google
import pytest_asyncio
from google.auth import compute_engine
from google.cloud import secretmanager, storage


#### Define Utility Functions
def get_env_var(key: str) -> str:
"""Gets environment variables."""
value = os.environ.get(key)
if value is None:
raise ValueError(f"Must set env var {key}")
return value


def access_secret_version(
project_id: str, secret_id: str, version_id: str = "latest"
) -> str:
"""Accesses the payload of a given secret version from Secret Manager."""
client = secretmanager.SecretManagerServiceClient()
name = f"projects/{project_id}/secrets/{secret_id}/versions/{version_id}"
response = client.access_secret_version(request={"name": name})
return response.payload.data.decode("UTF-8")


def create_tmpfile(content: str) -> str:
"""Creates a temporary file with the given content."""
with tempfile.NamedTemporaryFile(delete=False, mode="w") as tmpfile:
tmpfile.write(content)
return tmpfile.name


def download_blob(
bucket_name: str, source_blob_name: str, destination_file_name: str
) -> None:
"""Downloads a blob from a GCS bucket."""
storage_client = storage.Client()

bucket = storage_client.bucket(bucket_name)
blob = bucket.blob(source_blob_name)
blob.download_to_filename(destination_file_name)

print(f"Blob {source_blob_name} downloaded to {destination_file_name}.")


def get_toolbox_binary_url(toolbox_version: str) -> str:
"""Constructs the GCS path to the toolbox binary."""
os_system = platform.system().lower()
arch = (
"arm64" if os_system == "darwin" and platform.machine() == "arm64" else "amd64"
)
return f"v{toolbox_version}/{os_system}/{arch}/toolbox"


def get_auth_token(client_id: str) -> str:
"""Retrieves an authentication token"""
request = google.auth.transport.requests.Request()
credentials = compute_engine.IDTokenCredentials(
request=request,
target_audience=client_id,
use_metadata_identity_endpoint=True,
)
if not credentials.valid:
credentials.refresh(request)
return credentials.token


#### Define Fixtures
@pytest_asyncio.fixture(scope="session")
def project_id() -> str:
return get_env_var("GOOGLE_CLOUD_PROJECT")


@pytest_asyncio.fixture(scope="session")
def toolbox_version() -> str:
return get_env_var("TOOLBOX_VERSION")


@pytest_asyncio.fixture(scope="session")
def tools_file_path(project_id: str) -> Generator[str]:
"""Provides a temporary file path containing the tools manifest."""
tools_manifest = access_secret_version(
project_id=project_id,
secret_id="sdk_testing_tools",
version_id=os.environ.get("TOOLBOX_MANIFEST_VERSION", "latest"),
)
tools_file_path = create_tmpfile(tools_manifest)
yield tools_file_path
os.remove(tools_file_path)


@pytest_asyncio.fixture(scope="session")
def auth_token1(project_id: str) -> str:
client_id = access_secret_version(
project_id=project_id, secret_id="sdk_testing_client1"
)
return get_auth_token(client_id)


@pytest_asyncio.fixture(scope="session")
def auth_token2(project_id: str) -> str:
client_id = access_secret_version(
project_id=project_id, secret_id="sdk_testing_client2"
)
return get_auth_token(client_id)


@pytest_asyncio.fixture(scope="session")
def toolbox_server(toolbox_version: str, tools_file_path: str) -> Generator[None]:
"""Starts the toolbox server as a subprocess."""
print("Downloading toolbox binary from gcs bucket...")
source_blob_name = get_toolbox_binary_url(toolbox_version)
download_blob("genai-toolbox", source_blob_name, "toolbox")

print("Toolbox binary downloaded successfully.")
try:
print("Opening toolbox server process...")
# Make toolbox executable
os.chmod("toolbox", 0o700)
# Run toolbox binary
toolbox_server = subprocess.Popen(
["./toolbox", "--tools-file", tools_file_path]
)

# Wait for server to start
# Retry logic with a timeout
for _ in range(5): # retries
time.sleep(2)
print("Checking if toolbox is successfully started...")
if toolbox_server.poll() is None:
print("Toolbox server started successfully.")
break
else:
raise RuntimeError("Toolbox server failed to start after 5 retries.")
except subprocess.CalledProcessError as e:
print(e.stderr.decode("utf-8"))
print(e.stdout.decode("utf-8"))
raise RuntimeError(f"{e}\n\n{e.stderr.decode('utf-8')}") from e
yield

# Clean up toolbox server
toolbox_server.terminate()
toolbox_server.wait(timeout=5)
Loading