Skip to content
Merged
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
47 changes: 33 additions & 14 deletions agent_starter_pack/cli/commands/register_gemini_enterprise.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,10 @@
from rich.console import Console

from agent_starter_pack.cli.utils.command import run_gcloud_command
from agent_starter_pack.cli.utils.gcp import (
get_user_agent,
get_x_goog_api_client_header,
)

# TOML parser - use standard library for Python 3.11+, fallback to tomli
if sys.version_info >= (3, 11):
Expand Down Expand Up @@ -298,6 +302,32 @@ def get_access_token() -> str:
raise RuntimeError("Failed to get access token") from e


def _build_api_headers(
access_token: str,
project_id: str,
content_type: bool = False,
) -> dict[str, str]:
"""Build headers for Discovery Engine API requests with user-agent.

Args:
access_token: Google Cloud access token
project_id: GCP project ID or number for billing
content_type: Whether to include Content-Type header (for POST/PATCH)

Returns:
Headers dictionary
"""
headers = {
"Authorization": f"Bearer {access_token}",
"x-goog-user-project": project_id,
"User-Agent": get_user_agent(),
"x-goog-api-client": get_x_goog_api_client_header(),
}
if content_type:
headers["Content-Type"] = "application/json"
return headers


def get_identity_token() -> str:
"""Get Google Cloud identity token.

Expand Down Expand Up @@ -604,10 +634,7 @@ def list_gemini_enterprise_apps(
f"{base_endpoint}/v1alpha/projects/{project_number}/"
f"locations/{location}/collections/default_collection/engines"
)
headers = {
"Authorization": f"Bearer {access_token}",
"x-goog-user-project": project_number,
}
headers = _build_api_headers(access_token, project_number)

response = requests.get(url, headers=headers, timeout=30)
response.raise_for_status()
Expand Down Expand Up @@ -937,11 +964,7 @@ def register_a2a_agent(
f"locations/{as_location}/collections/{collection}/engines/{engine_id}/"
"assistants/default_assistant/agents"
)
headers = {
"Authorization": f"Bearer {access_token}",
"Content-Type": "application/json",
"x-goog-user-project": project_id,
}
headers = _build_api_headers(access_token, project_id, content_type=True)

# Build payload with A2A agent definition
payload = {
Expand Down Expand Up @@ -1106,11 +1129,7 @@ def register_agent(
)

# Request headers
headers = {
"Authorization": f"Bearer {access_token}",
"Content-Type": "application/json",
"x-goog-user-project": project_id,
}
headers = _build_api_headers(access_token, project_id, content_type=True)

# Request body
payload: dict = {
Expand Down
44 changes: 21 additions & 23 deletions agent_starter_pack/cli/utils/cicd.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@
from rich.prompt import IntPrompt, Prompt

from agent_starter_pack.cli.utils.command import get_gcloud_cmd
from agent_starter_pack.cli.utils.gcp import get_project_number

console = Console()

Expand Down Expand Up @@ -149,18 +150,7 @@ def create_github_connection(

# Get the Cloud Build service account and grant permissions with retry logic
try:
project_number_result = run_command(
[
"gcloud",
"projects",
"describe",
project_id,
"--format=value(projectNumber)",
],
capture_output=True,
check=True,
)
project_number = project_number_result.stdout.strip()
project_number = get_project_number(project_id)
cloud_build_sa = (
f"service-{project_number}@gcp-sa-cloudbuild.iam.gserviceaccount.com"
)
Expand Down Expand Up @@ -219,7 +209,14 @@ def create_github_connection(
"⏳ Waiting for IAM permissions to propagate (this typically takes 5-10 seconds)..."
)
time.sleep(10) # Give IAM time to propagate before proceeding
except subprocess.CalledProcessError as e:
except (PermissionError, ValueError) as e:
console.print(
f"⚠️ Could not setup Cloud Build service account: {e}", style="yellow"
)
console.print(
"You may need to manually grant the permissions if the connection creation fails."
)
except Exception as e:
console.print(
f"⚠️ Could not setup Cloud Build service account: {e}", style="yellow"
)
Expand Down Expand Up @@ -477,16 +474,7 @@ def ensure_apis_enabled(project_id: str, apis: list[str]) -> None:
capture_output=True,
)

project_number = run_command(
[
"gcloud",
"projects",
"describe",
project_id,
"--format=value(projectNumber)",
],
capture_output=True,
).stdout.strip()
project_number = get_project_number(project_id)

cloudbuild_sa = (
f"service-{project_number}@gcp-sa-cloudbuild.iam.gserviceaccount.com"
Expand All @@ -507,11 +495,21 @@ def ensure_apis_enabled(project_id: str, apis: list[str]) -> None:
)
console.print("✅ Permissions granted to Cloud Build service account")

except (PermissionError, ValueError) as e:
console.print(
f"❌ Failed to set up service account permissions: {e!s}", style="bold red"
)
raise
except subprocess.CalledProcessError as e:
console.print(
f"❌ Failed to set up service account permissions: {e!s}", style="bold red"
)
raise
except Exception as e:
console.print(
f"❌ Failed to set up service account permissions: {e!s}", style="bold red"
)
raise

# Add a small delay to allow API enablement and IAM changes to propagate
time.sleep(10)
Expand Down
54 changes: 49 additions & 5 deletions agent_starter_pack/cli/utils/gcp.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,9 @@
from agent_starter_pack.cli.utils.command import run_gcloud_command
from agent_starter_pack.cli.utils.version import PACKAGE_NAME, get_current_version

# API endpoint constants
RESOURCE_MANAGER_API_BASE = "https://cloudresourcemanager.googleapis.com"

# Lazy console - only create when needed
_console = None

Expand All @@ -49,16 +52,16 @@ def _get_console() -> Console:
)


def _get_user_agent(context: str | None = None) -> str:
def get_user_agent(context: str | None = None) -> str:
"""Returns a custom user agent string."""
version = get_current_version()
prefix = "ag" if context == "agent-garden" else ""
return f"{prefix}{version}-{PACKAGE_NAME}/{prefix}{version}-{PACKAGE_NAME}"


def _get_x_goog_api_client_header(context: str | None = None) -> str:
def get_x_goog_api_client_header(context: str | None = None) -> str:
"""Build x-goog-api-client header matching SDK format."""
user_agent = _get_user_agent(context)
user_agent = get_user_agent(context)
python_version = (
f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}"
)
Expand Down Expand Up @@ -112,8 +115,8 @@ def _test_vertex_connection(
Returns:
Tuple of (success, error_message)
"""
user_agent = _get_user_agent(context)
x_goog_api_client = _get_x_goog_api_client_header(context)
user_agent = get_user_agent(context)
x_goog_api_client = get_x_goog_api_client_header(context)

try:
response = requests.post(
Expand Down Expand Up @@ -268,3 +271,44 @@ def get_account() -> str:
):
raise Exception(_AUTH_ERROR_MESSAGE) from e
raise


def get_project_number(project_id: str) -> str:
"""Get project number from project ID using Resource Manager API.

Args:
project_id: GCP project ID

Returns:
Project number as string

Raises:
PermissionError: If access is denied to the project
ValueError: If the project is not found
requests.exceptions.HTTPError: For other API failures
"""
_, _, token = _get_credentials_and_token()

user_agent = get_user_agent()
x_goog_api_client = get_x_goog_api_client_header()

response = requests.get(
f"{RESOURCE_MANAGER_API_BASE}/v1/projects/{project_id}",
headers={
"Authorization": f"Bearer {token}",
"User-Agent": user_agent,
"x-goog-api-client": x_goog_api_client,
},
timeout=30,
)

if response.status_code == 403:
raise PermissionError(
f"Permission denied accessing project '{project_id}'. "
"Ensure you have the required permissions."
)
if response.status_code == 404:
raise ValueError(f"Project '{project_id}' not found.")

response.raise_for_status()
return response.json()["projectNumber"]
1 change: 0 additions & 1 deletion docs/guide/development-guide.md
Original file line number Diff line number Diff line change
Expand Up @@ -92,7 +92,6 @@ go get <package> # Add dependency
go mod tidy # Clean up dependencies
```
:::
:::

> Note: The specific UI playground launched by `make playground` depends on the agent template you selected during creation.

Expand Down
5 changes: 5 additions & 0 deletions docs/package-lock.json

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.