From 99b308078f8489ff3520cd7be873f76c6d42d13b Mon Sep 17 00:00:00 2001 From: Abhishek Singh Date: Mon, 31 Aug 2026 15:06:43 -0700 Subject: [PATCH 1/2] Added support for Select AI A2A Gateway --- .pre-commit-config.yaml | 1 + gcloud/README.md | 172 +-------- gcloud/gateway/README.md | 145 ++++++++ gcloud/gateway/cloudbuild.yaml | 220 ++++++++++++ gcloud/gateway/deploy.sh | 236 ++++++++++++ gcloud/gateway/gke/consul.yaml | 50 +++ gcloud/gateway/gke/namespace.yaml | 4 + gcloud/gateway/gke/worker-mtls.yaml | 82 +++++ gcloud/gateway/gke/worker.yaml | 52 +++ gcloud/standalone/README.md | 159 +++++++++ gcloud/{ => standalone}/cloudbuild.yaml | 4 +- gcloud/{ => standalone}/deploy.sh | 6 +- pyproject.toml | 3 + src/select_ai/agent/a2a/__init__.py | 30 +- src/select_ai/agent/a2a/a2ui.py | 33 ++ src/select_ai/agent/a2a/forms.py | 117 ++++++ src/select_ai/agent/a2a/gateway.py | 231 ++++++++++++ src/select_ai/agent/a2a/models.py | 71 ++++ src/select_ai/agent/a2a/results.py | 58 +++ src/select_ai/agent/a2a/server.py | 48 +-- src/select_ai/agent/a2a/worker.py | 436 +++++++++++++++++++++++ src/select_ai/agent/a2a/worker_client.py | 174 +++++++++ src/select_ai/cli/a2a.py | 151 ++++++++ tests/a2a/test_agent_card.py | 55 +++ tests/a2a/test_worker_runtime.py | 382 ++++++++++++++++++++ 25 files changed, 2720 insertions(+), 200 deletions(-) create mode 100644 gcloud/gateway/README.md create mode 100644 gcloud/gateway/cloudbuild.yaml create mode 100755 gcloud/gateway/deploy.sh create mode 100644 gcloud/gateway/gke/consul.yaml create mode 100644 gcloud/gateway/gke/namespace.yaml create mode 100644 gcloud/gateway/gke/worker-mtls.yaml create mode 100644 gcloud/gateway/gke/worker.yaml create mode 100644 gcloud/standalone/README.md rename gcloud/{ => standalone}/cloudbuild.yaml (82%) rename gcloud/{ => standalone}/deploy.sh (98%) create mode 100644 src/select_ai/agent/a2a/a2ui.py create mode 100644 src/select_ai/agent/a2a/forms.py create mode 100644 src/select_ai/agent/a2a/gateway.py create mode 100644 src/select_ai/agent/a2a/models.py create mode 100644 src/select_ai/agent/a2a/results.py create mode 100644 src/select_ai/agent/a2a/worker.py create mode 100644 src/select_ai/agent/a2a/worker_client.py create mode 100644 tests/a2a/test_worker_runtime.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 3f3609a..a0fea1a 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -5,6 +5,7 @@ repos: - id: trailing-whitespace - id: end-of-file-fixer - id: check-yaml + args: [--allow-multiple-documents] - id: check-added-large-files args: ['--maxkb=2000'] - repo: https://github.com/psf/black-pre-commit-mirror diff --git a/gcloud/README.md b/gcloud/README.md index e075515..cf52356 100644 --- a/gcloud/README.md +++ b/gcloud/README.md @@ -1,159 +1,19 @@ -# Deploy the Select AI A2A server to Google Cloud + # Google Cloud deployment modes -`gcloud/deploy.sh` builds or selects a Select AI container image, creates or -updates a private Cloud Run service, and configures its database secrets. Run -it on a machine with the [Google Cloud CLI](https://cloud.google.com/sdk/docs/install) -installed and authenticated to the target project. +Select AI supports two distinct Google Cloud deployment modes. Choose based on +whether the database and Select AI team are known at deployment time or must +be selected dynamically by each user. -## IAM permissions - -The scripts use the active `gcloud` identity. They do not elevate its access. - -### Deployer (the active gcloud identity) - -| Operation | Required permissions | -| --- | --- | -| Inspect and create the Artifact Registry repository | `artifactregistry.repositories.get`, `artifactregistry.repositories.create` | -| Inspect and create the default runtime service account | `iam.serviceAccounts.get`, `iam.serviceAccounts.create` | -| Deploy or update Cloud Run | `run.services.create`, `run.services.update`, `run.services.get`, `run.operations.get`; `iam.serviceAccounts.actAs` on the runtime service account; `artifactregistry.repositories.downloadArtifacts` on the image repository | -| With `--build`, upload local source, submit, and wait for a build | `storage.buckets.get`, `storage.objects.create` on the configured source-staging bucket; `cloudbuild.builds.create`, `cloudbuild.builds.get`, `serviceusage.services.use` | -| Inspect, create, and add versions to database or wallet secrets | `secretmanager.secrets.get`, `secretmanager.secrets.create`, `secretmanager.versions.add` | -| Grant the runtime account access to those secrets | `secretmanager.secrets.getIamPolicy`, `secretmanager.secrets.setIamPolicy` | -| Grant Gemini Enterprise and the active gcloud identity access to the service | `run.services.getIamPolicy`, `run.services.setIamPolicy` | -| Obtain the project number | `resourcemanager.projects.get` | - -### Runtime service account - -| Operation | Required permissions | -| --- | --- | -| Read database and wallet secrets while serving requests | `secretmanager.versions.access` | - -### Other service identities - -| Principal | Operation | Required permissions | +| | Standalone | Dynamic gateway | | --- | --- | --- | -| Cloud Build execution service account | With `--build`, push the built image | `artifactregistry.repositories.uploadArtifacts` | -| Gemini Enterprise service agent | Invoke the private Cloud Run service | `run.routes.invoke` | -| Active gcloud identity | Fetch the Agent Card after deployment | `run.routes.invoke` | - -The source-staging bucket is Cloud Build's default unless a custom bucket is -configured. Cloud Build also needs access to its build-log destination; the -default same-project build account has that access. If your organization uses -a custom build service account, source bucket, or log bucket, its administrator -must grant the equivalent Cloud Storage permissions on those resources. - -Google Cloud references: [Service Usage access control](https://cloud.google.com/service-usage/docs/access-control), [Cloud Run deployment permissions](https://cloud.google.com/run/docs/reference/iam/roles), [Secret Manager access control](https://cloud.google.com/secret-manager/docs/access-control), [Artifact Registry roles](https://cloud.google.com/artifact-registry/docs/access-control), and [Cloud Build roles](https://cloud.google.com/build/docs/iam-roles-permissions). - -The rows that set IAM policy are administrative mutations. They are present -because `deploy.sh` creates and rotates secrets and configures private-service -invocation. If your customer deployment identity must not change IAM, provision -the secrets and the `secretmanager.versions.access`/`run.routes.invoke` -permissions beforehand, then -remove those policy-setting commands from the deployment workflow. - -## Prerequisite: enable project APIs once - -An administrator must enable these APIs once for the project: - -```bash -gcloud services enable \ - run.googleapis.com \ - cloudbuild.googleapis.com \ - artifactregistry.googleapis.com \ - secretmanager.googleapis.com \ - discoveryengine.googleapis.com \ - --project PROJECT_ID -``` - -## Deploy (and update) the A2A server - -```bash -gcloud/deploy.sh --build -``` - -On the first deployment, the script prompts for the ADB user, password, and -connect descriptor. It stores them in Secret Manager under names based on the -Cloud Run service, and grants only the runtime service account access. The -container receives the values as `SELECT_AI_USER`, `SELECT_AI_PASSWORD`, and -`SELECT_AI_DB_CONNECT_STRING`; they are never placed in the image or source -tree. - -### Optional: Autonomous Database mTLS wallet - -The Select AI SDK already supports `wallet_location` and `wallet_password`. -For Cloud Run, pass the path to the downloaded Autonomous Database wallet ZIP -on the first deployment (or when replacing it): - -```bash -gcloud/deploy.sh --wallet-archive /path/to/Wallet_database.zip -``` - -The script prompts for the wallet password, stores the ZIP and password as -service-specific Secret Manager secrets, and grants access only to the runtime -service account. Cloud Run mounts the ZIP read-only; its A2A launcher expands it -into ephemeral `/tmp` storage before starting the SDK, verifies it contains -`ewallet.pem`, and sets `SELECT_AI_WALLET_LOCATION` to that file's directory. -Do not commit the wallet ZIP or put its contents in the image. - -Later deploys reuse the wallet. To replace it, pass `--wallet-archive` again. - -The first deployment needs `--build` (or an explicit `--image-uri`). Later -deployments reuse the image already deployed to the service, so changing Cloud -Run configuration or secrets does not create another image. The command -deploys private Cloud Run, sets the final public URL in the Agent Card, grants your active -gcloud identity and Gemini Enterprise Discovery Engine service agent the -`run.routes.invoke` permission for this Cloud Run service. - -The default Cloud Run service is `oracle-a2a-agent`. Its default Agent Team, -installed in Oracle Database, is `ORACLE_AI_DATABASE_AGENT`. Override either -with explicit options: - -```bash -gcloud/deploy.sh --service sales-analyst-a2a --a2a-team SALES_ANALYST -``` - -Use a distinct `--service` value for each A2A team. Each service gets distinct Secret -Manager secret names by default, so credentials remain attached to that A2A -server. - -`--max-instances` controls the number of Cloud Run containers. Each container -can use up to 10 Oracle connections by default; change that limit with -`--pool-max-size`, for example `gcloud/deploy.sh --pool-max-size 20`. - -### Update the Select AI SDK or this repository - -Update the checkout (or modify its dependency version), then explicitly build -and deploy the new image: - -```bash -git pull -gcloud/deploy.sh --build -``` - -`--build` creates a freshly tagged image from the current source; without it, -the existing image is reused. Existing database secrets are reused without -prompting. To rotate the ADB credentials, explicitly request it: - -```bash -gcloud/deploy.sh --rotate-db-credentials -``` - -### What `cloudbuild.yaml` does - -`gcloud/deploy.sh --build` uses `gcloud/cloudbuild.yaml` to tell Cloud Build to build -`docker/Dockerfile` and push it to Artifact Registry. It is build configuration, -not a command you run. The build context is the repository root, so the image -can install the Select AI source from `pyproject.toml` and `src/`. - -### Cloud Build upload contents - -Before the build starts, `gcloud builds submit` archives and uploads the -repository root. The root `.gcloudignore` excludes local virtual environments, -generated documentation, test data, caches, credentials, and Git metadata. -Keep `src/`, `pyproject.toml`, `docker/`, and `gcloud/` in the upload; they are -required to build the image. If the upload is unexpectedly large, check local -directories against `.gcloudignore` before running `--build` again. - -After a successful deployment, the script prints the A2A Agent Card JSON. -Paste that JSON into Gemini Enterprise to register the private service. The -required Gemini Enterprise invocation permission has already been added. +| Database and team | Fixed at deployment time | Chosen at runtime for each user session | +| Public A2A service | One service for one configured team | One gateway that presents an A2UI connection form | +| Users | All requests use the deployed database identity | Any permitted user can connect to a reachable Oracle database and Select AI team | +| Architecture | One Cloud Run service | Cloud Run gateway, plus Consul and worker replicas in GKE | +| Session isolation | Shared service database pool | One child process and async pool per active user session | +| Main benefit | Simple, predictable deployment | Dynamic, multi-database and multi-team access from one A2A endpoint | +| Operational cost | Low | Higher: GKE workers, Consul, routing, TTL, and session capacity | + +Use [standalone](standalone/README.md) when a service should expose one known +database team. Use [gateway](gateway/README.md) when users must dynamically +choose their database connection and team. diff --git a/gcloud/gateway/README.md b/gcloud/gateway/README.md new file mode 100644 index 0000000..a992159 --- /dev/null +++ b/gcloud/gateway/README.md @@ -0,0 +1,145 @@ +# Dynamic gateway deployment + +Dynamic gateway mode exposes one public A2A endpoint where each user selects +an Oracle database connection and Select AI team through the A2UI connection +form. + +```text +A2A client + | +Cloud Run gateway + | + +-- private VPC --> Consul in GKE + | + +-- private VPC --> worker replicas in GKE + | + session child process + | + Oracle Database +``` + +The gateway is the only public A2A application. Consul and workers are a GKE +clustered service: Consul selects a worker for each new dynamic session, and +the chosen worker retains that session's process and Oracle conversation. + +## Deploy the complete stack + +Run this from the repository root: + +```bash +gcloud/gateway/deploy.sh --project PROJECT_ID +``` + +The script creates the Artifact Registry repository and GKE Autopilot cluster +when they do not already exist. Cloud Build then: + +1. builds the existing `docker/Dockerfile` image once; +2. deploys the GKE namespace and internal Consul service; +3. deploys the requested number of GKE worker replicas using + `select-ai a2a worker`; +4. deploys the same image to Cloud Run using `select-ai a2a gateway`; +5. sets the final Cloud Run URL in `AGENT_URL` for the Agent Card. + +Common options: + +```bash +gcloud/gateway/deploy.sh \ + --project PROJECT_ID \ + --region us-central1 \ + --cluster select-ai-a2a-gateway \ + --worker-replicas 3 \ + --network default \ + --subnet default +``` + +The Cloud Run gateway uses direct VPC egress to reach the internal Consul load +balancer and GKE worker pod addresses. The default one-instance gateway limit +is intentional: gateway A2A task and context/session state is currently in +memory. Workers, rather than the gateway, provide the clustered capacity for +dynamic sessions. + +`cloudbuild.yaml` is the complete build and deployment workflow. It supplies +the generated image and Consul endpoint values to the Cloud Run gateway at +deployment time. + +## Optional worker mTLS test mode + +Local testing does not use mTLS. The default GCloud deployment also keeps the +current private-VPC HTTP worker transport. + +For a short-lived GCloud mTLS test, add a private DNS suffix: + +```bash +gcloud/gateway/deploy.sh \ + --project PROJECT_ID \ + --enable-worker-mtls \ + --worker-domain workers.select-ai.internal \ + --mtls-cert-validity-days 365 +``` + +This mode is intentionally self-contained and is not a production PKI design. +The deployment script generates an ephemeral CA and leaf certificates valid for +365 days by default, then removes the CA private key from its restricted +temporary directory. It never prints or stores that key. The leaf material is +first stored in Google Secret Manager. Cloud Build then creates the Kubernetes +Secrets used by the workers. + +Set `--mtls-cert-validity-days DAYS` to choose the lifetime for both leaf +certificates: the gateway client certificate and the worker server certificate. +The CA is issued for one additional day. + +The first mTLS deployment creates these certificates. Later mTLS deployments +reuse them, including when changing `--worker-replicas`. To deliberately +replace the CA and both leaf certificates, add `--rotate-worker-mtls`. Rotation +recreates the worker StatefulSet and ends active worker sessions. + +Workers run as a StatefulSet. Each worker registers its own private DNS name +with Consul, so the gateway still reaches the exact process that owns the +session. Cloud Build creates the corresponding private Cloud DNS A records. + +The certificate and DNS records are refreshed by rerunning the deployment. +Because this is a test mode, rerun it after a worker pod is recreated outside a +deployment; its pod IP can change. Do not use this mode for a long-lived +production stack—use a managed workload-certificate and DNS reconciliation +solution there. + +### Where the worker certificate files come from + +The certificate paths passed to `select-ai a2a worker` are files **inside each +GKE worker container**. They are mounted from Kubernetes Secrets, not files in +this repository or on the machine that runs `deploy.sh`. + +```text +deploy.sh + -> creates the test certificates and stores them in Secret Manager + -> Cloud Build reads the worker certificate, worker key, and CA certificate + -> Cloud Build creates Kubernetes Secrets in select-ai-gateway + -> GKE mounts those Secrets read-only in every worker container +``` + +The worker pod receives these mounts: + +| Container file | Kubernetes Secret | Secret key | Used for | +| --- | --- | --- | --- | +| `/var/run/select-ai-mtls/tls.crt` | `select-ai-worker-server-tls` | `tls.crt` | worker HTTPS server certificate | +| `/var/run/select-ai-mtls/tls.key` | `select-ai-worker-server-tls` | `tls.key` | worker HTTPS private key | +| `/var/run/select-ai-mtls/gateway-ca.crt` | `select-ai-gateway-client-ca` | `ca.crt` | validates the gateway client certificate | + +GKE does not read Google Secret Manager at worker request time. Cloud Build +copies the required material into Kubernetes Secrets during deployment, and the +worker certificate files are mounted from those Kubernetes Secrets as read-only +files. For example, to inspect the paths in a running worker (without printing +their contents): + +```bash +kubectl -n select-ai-gateway exec select-ai-worker-0 -- \ + ls -l /var/run/select-ai-mtls/ +``` + +The Cloud Run gateway uses a separate direct mount from Google Secret Manager +for its CA, client certificate, and client key. + +The identity that submits Cloud Build needs permission to use GKE, Cloud Run, +Cloud DNS, and Secret Manager. Cloud Build also needs the corresponding GKE, +Cloud Run, Cloud DNS, and Secret Manager permissions because it deploys the +workers, refreshes the private DNS records, and mounts the gateway secrets. diff --git a/gcloud/gateway/cloudbuild.yaml b/gcloud/gateway/cloudbuild.yaml new file mode 100644 index 0000000..69958b7 --- /dev/null +++ b/gcloud/gateway/cloudbuild.yaml @@ -0,0 +1,220 @@ +# Build one Select AI image, deploy Consul and workers to GKE, then deploy the +# public dynamic gateway to Cloud Run. +steps: + - name: gcr.io/cloud-builders/docker + args: + - build + - --file + - docker/Dockerfile + - --tag + - ${_REGION}-docker.pkg.dev/$PROJECT_ID/${_REPOSITORY}/select-ai:${_IMAGE_TAG} + - . + + # GKE and Cloud Run consume the image during this build, so publish it before + # deploying either workload. The top-level images: field publishes only after + # every build step succeeds, which is too late for the worker rollout. + - name: gcr.io/cloud-builders/docker + args: + - push + - ${_REGION}-docker.pkg.dev/$PROJECT_ID/${_REPOSITORY}/select-ai:${_IMAGE_TAG} + + - name: gcr.io/cloud-builders/gke-deploy + args: + - run + - --filename=gcloud/gateway/gke/namespace.yaml + - --location=${_REGION} + - --cluster=${_CLUSTER} + + # gke-deploy writes generated manifests to output/. Cloud Build shares the + # source workspace across steps, so remove only that builder output before + # invoking it again. + - name: gcr.io/cloud-builders/gcloud + entrypoint: bash + args: + - -ceu + - rm -rf output + + - name: gcr.io/cloud-builders/kubectl + entrypoint: bash + env: + - CLOUDSDK_COMPUTE_REGION=${_REGION} + - CLOUDSDK_CONTAINER_CLUSTER=${_CLUSTER} + args: + - -ceu + - | + if [[ "${_ENABLE_WORKER_MTLS}" == "true" ]]; then + kubectl delete --namespace select-ai-gateway deployment/select-ai-worker --ignore-not-found + if [[ "${_ROTATE_WORKER_MTLS}" == "true" ]]; then + kubectl delete --namespace select-ai-gateway statefulset/select-ai-worker --ignore-not-found --wait=true + fi + else + kubectl delete --namespace select-ai-gateway statefulset/select-ai-worker --ignore-not-found + fi + + - name: gcr.io/cloud-builders/gcloud + entrypoint: bash + args: + - -ceu + - | + if [[ "${_ENABLE_WORKER_MTLS}" != "true" ]]; then + exit 0 + fi + gcloud secrets versions access latest --secret=select-ai-worker-mtls-cert > /workspace/worker.crt + gcloud secrets versions access latest --secret=select-ai-worker-mtls-key > /workspace/worker.key + gcloud secrets versions access latest --secret=select-ai-gateway-mtls-ca > /workspace/gateway-ca.crt + + - name: gcr.io/cloud-builders/kubectl + entrypoint: bash + env: + - CLOUDSDK_COMPUTE_REGION=${_REGION} + - CLOUDSDK_CONTAINER_CLUSTER=${_CLUSTER} + args: + - -ceu + - | + if [[ "${_ENABLE_WORKER_MTLS}" != "true" ]]; then + exit 0 + fi + kubectl -n select-ai-gateway create secret tls select-ai-worker-server-tls \ + --cert=/workspace/worker.crt --key=/workspace/worker.key \ + --dry-run=client -o yaml | kubectl apply -f - + kubectl -n select-ai-gateway create secret generic select-ai-gateway-client-ca \ + --from-file=ca.crt=/workspace/gateway-ca.crt \ + --dry-run=client -o yaml | kubectl apply -f - + + - name: gcr.io/cloud-builders/gke-deploy + args: + - run + - --filename=gcloud/gateway/gke/consul.yaml + - --location=${_REGION} + - --cluster=${_CLUSTER} + + - name: gcr.io/cloud-builders/kubectl + entrypoint: bash + env: + - CLOUDSDK_COMPUTE_REGION=${_REGION} + - CLOUDSDK_CONTAINER_CLUSTER=${_CLUSTER} + args: + - -ceu + - | + kubectl wait --namespace select-ai-gateway \ + --for=jsonpath='{.status.loadBalancer.ingress[0].ip}' \ + service/consul-server --timeout=300s + kubectl get --namespace select-ai-gateway service/consul-server \ + -o jsonpath='{.status.loadBalancer.ingress[0].ip}' \ + > /workspace/consul-ip + + - name: gcr.io/cloud-builders/gcloud + entrypoint: bash + args: + - -ceu + - | + image="${_REGION}-docker.pkg.dev/$PROJECT_ID/${_REPOSITORY}/select-ai:${_IMAGE_TAG}" + worker_manifest=gcloud/gateway/gke/worker.yaml + if [[ "${_ENABLE_WORKER_MTLS}" == "true" ]]; then + worker_manifest=gcloud/gateway/gke/worker-mtls.yaml + fi + sed -e "s|WORKER_IMAGE|$image|" \ + -e "s|WORKER_REPLICAS|${_WORKER_REPLICAS}|" \ + -e "s|WORKER_DOMAIN|${_WORKER_DOMAIN}|" \ + "$worker_manifest" > /workspace/worker.yaml + + - name: gcr.io/cloud-builders/gcloud + entrypoint: bash + args: + - -ceu + - rm -rf output + + - name: gcr.io/cloud-builders/gke-deploy + args: + - run + - --filename=/workspace/worker.yaml + - --location=${_REGION} + - --cluster=${_CLUSTER} + + - name: gcr.io/cloud-builders/kubectl + entrypoint: bash + env: + - CLOUDSDK_COMPUTE_REGION=${_REGION} + - CLOUDSDK_CONTAINER_CLUSTER=${_CLUSTER} + args: + - -ceu + - | + if [[ "${_ENABLE_WORKER_MTLS}" == "true" ]]; then + kubectl rollout status --namespace select-ai-gateway statefulset/select-ai-worker --timeout=300s + kubectl get pods --namespace select-ai-gateway -l app=select-ai-worker \ + -o jsonpath='{range .items[*]}{.metadata.name}={.status.podIP}{"\n"}{end}' \ + > /workspace/worker-pods + else + kubectl rollout status --namespace select-ai-gateway deployment/select-ai-worker --timeout=300s + fi + + - name: gcr.io/cloud-builders/gcloud + entrypoint: bash + args: + - -ceu + - | + if [[ "${_ENABLE_WORKER_MTLS}" != "true" ]]; then + exit 0 + fi + while IFS== read -r pod_name pod_ip; do + record_name="${pod_name}.${_WORKER_DOMAIN}." + gcloud dns record-sets delete "$record_name" --type=A \ + --zone="${_WORKER_DNS_ZONE}" --quiet || true + gcloud dns record-sets create "$record_name" --type=A --ttl=30 \ + --rrdatas="$pod_ip" --zone="${_WORKER_DNS_ZONE}" + done < /workspace/worker-pods + + - name: gcr.io/cloud-builders/gcloud + entrypoint: bash + args: + - -ceu + - | + consul_ip=$(cat /workspace/consul-ip) + image="${_REGION}-docker.pkg.dev/$PROJECT_ID/${_REPOSITORY}/select-ai:${_IMAGE_TAG}" + mtls_args=() + mtls_env="" + if [[ "${_ENABLE_WORKER_MTLS}" == "true" ]]; then + mtls_args=( + "--update-secrets=/var/run/select-ai-mtls-ca/ca.pem=select-ai-gateway-mtls-ca:latest,/var/run/select-ai-mtls-cert/tls.crt=select-ai-gateway-mtls-cert:latest,/var/run/select-ai-mtls-key/tls.key=select-ai-gateway-mtls-key:latest" + ) + mtls_env=",WORKER_TLS_CA_FILE=/var/run/select-ai-mtls-ca/ca.pem,WORKER_TLS_CERT_FILE=/var/run/select-ai-mtls-cert/tls.crt,WORKER_TLS_KEY_FILE=/var/run/select-ai-mtls-key/tls.key" + fi + gcloud run deploy "${_GATEWAY_SERVICE}" \ + --image="$image" \ + --region="${_REGION}" \ + --no-allow-unauthenticated \ + --network="${_NETWORK}" \ + --subnet="${_SUBNET}" \ + --vpc-egress=all-traffic \ + --max-instances=1 \ + --min-instances=1 \ + --port=8080 \ + --command=select-ai \ + --args=a2a,gateway \ + --set-env-vars="AGENT_URL=https://pending.invalid,CONSUL_HTTP_URL=http://$consul_ip:8500,WORKER_SERVICE=select-ai-worker,SESSION_TTL_SECONDS=900$mtls_env" \ + "${mtls_args[@]}" + gateway_url=$(gcloud run services describe "${_GATEWAY_SERVICE}" \ + --region="${_REGION}" --format='value(status.url)') + gcloud run services update "${_GATEWAY_SERVICE}" \ + --region="${_REGION}" \ + --update-env-vars="AGENT_URL=$gateway_url" + +images: + - ${_REGION}-docker.pkg.dev/$PROJECT_ID/${_REPOSITORY}/select-ai:${_IMAGE_TAG} + +substitutions: + _REGION: us-central1 + _CLUSTER: select-ai-a2a-gateway + _REPOSITORY: select-ai + _IMAGE_TAG: dev + _GATEWAY_SERVICE: select-ai-a2a-gateway + _NETWORK: default + _SUBNET: default + _WORKER_REPLICAS: "2" + _ENABLE_WORKER_MTLS: "false" + _ROTATE_WORKER_MTLS: "false" + _WORKER_DOMAIN: "" + _WORKER_DNS_ZONE: select-ai-workers + +options: + logging: CLOUD_LOGGING_ONLY diff --git a/gcloud/gateway/deploy.sh b/gcloud/gateway/deploy.sh new file mode 100755 index 0000000..1923940 --- /dev/null +++ b/gcloud/gateway/deploy.sh @@ -0,0 +1,236 @@ +#!/usr/bin/env bash + +# ----------------------------------------------------------------------------- +# Copyright (c) 2026, Oracle and/or its affiliates. +# +# Licensed under the Universal Permissive License v 1.0 as shown at +# https://oss.oracle.com/licenses/upl. +# ----------------------------------------------------------------------------- + +# Build and deploy the dynamic Select AI gateway stack: Cloud Run gateway plus +# Consul and worker replicas in GKE. + +set -euo pipefail + +usage() { + cat <<'EOF' +Usage: gcloud/gateway/deploy.sh [options] + +Options: + --project PROJECT Google Cloud project (defaults to gcloud config) + --region REGION Region for GKE, Cloud Run, and Artifact Registry (default: us-central1) + --cluster NAME GKE Autopilot cluster name (default: select-ai-a2a-gateway) + --repository NAME Artifact Registry Docker repository (default: select-ai) + --gateway-service NAME Cloud Run gateway service name (default: select-ai-a2a-gateway) + --network NAME VPC network for Cloud Run direct VPC egress (default: default) + --subnet NAME VPC subnet for Cloud Run direct VPC egress (default: default) + --worker-replicas COUNT GKE worker replica count (default: 2) + --enable-worker-mtls Use ephemeral mTLS certificates for gateway-to-worker calls + --rotate-worker-mtls Replace the existing worker mTLS CA and certificates + --worker-domain DOMAIN Private DNS suffix for workers; required with --enable-worker-mtls + --worker-dns-zone NAME Private Cloud DNS zone name (default: select-ai-workers) + --mtls-cert-validity-days DAYS + Gateway and worker certificate lifetime (default: 365) + --image-tag TAG Image tag (default: git SHA plus UTC timestamp) + -h, --help Show this help +EOF +} + +repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +project_id="" +region="us-central1" +cluster="select-ai-a2a-gateway" +repository="select-ai" +gateway_service="select-ai-a2a-gateway" +network="default" +subnet="default" +worker_replicas="2" +enable_worker_mtls="false" +rotate_worker_mtls="false" +worker_domain="" +worker_dns_zone="select-ai-workers" +mtls_cert_validity_days="365" +image_tag="" + +while [[ $# -gt 0 ]]; do + case "$1" in + --project) project_id="${2:?--project requires a value}"; shift 2 ;; + --region) region="${2:?--region requires a value}"; shift 2 ;; + --cluster) cluster="${2:?--cluster requires a value}"; shift 2 ;; + --repository) repository="${2:?--repository requires a value}"; shift 2 ;; + --gateway-service) gateway_service="${2:?--gateway-service requires a value}"; shift 2 ;; + --network) network="${2:?--network requires a value}"; shift 2 ;; + --subnet) subnet="${2:?--subnet requires a value}"; shift 2 ;; + --worker-replicas) worker_replicas="${2:?--worker-replicas requires a value}"; shift 2 ;; + --enable-worker-mtls) enable_worker_mtls="true"; shift ;; + --rotate-worker-mtls) rotate_worker_mtls="true"; shift ;; + --worker-domain) worker_domain="${2:?--worker-domain requires a value}"; shift 2 ;; + --worker-dns-zone) worker_dns_zone="${2:?--worker-dns-zone requires a value}"; shift 2 ;; + --mtls-cert-validity-days) mtls_cert_validity_days="${2:?--mtls-cert-validity-days requires a value}"; shift 2 ;; + --image-tag) image_tag="${2:?--image-tag requires a value}"; shift 2 ;; + -h|--help) usage; exit 0 ;; + *) echo "Unknown option: $1" >&2; usage >&2; exit 2 ;; + esac +done + +if [[ -z "$project_id" ]]; then + project_id="$(gcloud config get-value project 2>/dev/null || true)" +fi +if [[ -z "$project_id" || "$project_id" == "(unset)" ]]; then + echo "Pass --project or configure one with: gcloud config set project PROJECT_ID" >&2 + exit 2 +fi +if ! [[ "$worker_replicas" =~ ^[1-9][0-9]*$ ]]; then + echo "--worker-replicas must be a positive integer." >&2 + exit 2 +fi +if [[ "$enable_worker_mtls" == "true" && -z "$worker_domain" ]]; then + echo "--worker-domain is required with --enable-worker-mtls." >&2 + exit 2 +fi +if [[ "$rotate_worker_mtls" == "true" && "$enable_worker_mtls" != "true" ]]; then + echo "--rotate-worker-mtls requires --enable-worker-mtls." >&2 + exit 2 +fi +if ! [[ "$mtls_cert_validity_days" =~ ^[1-9][0-9]*$ ]]; then + echo "--mtls-cert-validity-days must be a positive integer." >&2 + exit 2 +fi + +image_tag="${image_tag:-$(git -C "$repo_root" rev-parse --short HEAD)-$(date -u +%Y%m%d%H%M%S)}" + +gcloud services enable \ + artifactregistry.googleapis.com \ + cloudbuild.googleapis.com \ + compute.googleapis.com \ + container.googleapis.com \ + discoveryengine.googleapis.com \ + dns.googleapis.com \ + run.googleapis.com \ + secretmanager.googleapis.com \ + --project="$project_id" + +if ! gcloud artifacts repositories describe "$repository" \ + --location="$region" --project="$project_id" >/dev/null 2>&1; then + gcloud artifacts repositories create "$repository" \ + --repository-format=docker --location="$region" --project="$project_id" +fi + +if ! gcloud container clusters describe "$cluster" \ + --location="$region" --project="$project_id" >/dev/null 2>&1; then + gcloud container clusters create-auto "$cluster" \ + --location="$region" --network="$network" --project="$project_id" +fi + +cleanup_mtls() { + [[ -n "${mtls_dir:-}" ]] && rm -rf "$mtls_dir" +} + +if [[ "$enable_worker_mtls" == "true" ]]; then + create_mtls_material="$rotate_worker_mtls" + for secret_name in \ + select-ai-gateway-mtls-ca \ + select-ai-gateway-mtls-cert \ + select-ai-gateway-mtls-key \ + select-ai-worker-mtls-cert \ + select-ai-worker-mtls-key; do + if ! gcloud secrets versions access latest --secret="$secret_name" \ + --project="$project_id" >/dev/null 2>&1; then + create_mtls_material="true" + fi + done + if [[ "$create_mtls_material" == "true" ]]; then + trap cleanup_mtls EXIT + mtls_dir="$(mktemp -d)" + umask 077 + ca_validity_days="$((mtls_cert_validity_days + 1))" + openssl req -x509 -newkey rsa:2048 -nodes -days "$ca_validity_days" \ + -keyout "$mtls_dir/ca.key" -out "$mtls_dir/ca.crt" \ + -subj "/CN=select-ai ephemeral worker CA" >/dev/null 2>&1 + openssl req -newkey rsa:2048 -nodes \ + -keyout "$mtls_dir/worker.key" -out "$mtls_dir/worker.csr" \ + -subj "/CN=*.$worker_domain" >/dev/null 2>&1 + printf 'subjectAltName=DNS:*.%s\nextendedKeyUsage=serverAuth\n' "$worker_domain" \ + > "$mtls_dir/worker.ext" + openssl x509 -req -days "$mtls_cert_validity_days" -in "$mtls_dir/worker.csr" \ + -CA "$mtls_dir/ca.crt" -CAkey "$mtls_dir/ca.key" -CAcreateserial \ + -out "$mtls_dir/worker.crt" -extfile "$mtls_dir/worker.ext" >/dev/null 2>&1 + openssl req -newkey rsa:2048 -nodes \ + -keyout "$mtls_dir/gateway.key" -out "$mtls_dir/gateway.csr" \ + -subj "/CN=select-ai-gateway" >/dev/null 2>&1 + printf 'extendedKeyUsage=clientAuth\n' > "$mtls_dir/gateway.ext" + openssl x509 -req -days "$mtls_cert_validity_days" -in "$mtls_dir/gateway.csr" \ + -CA "$mtls_dir/ca.crt" -CAkey "$mtls_dir/ca.key" -CAcreateserial \ + -out "$mtls_dir/gateway.crt" -extfile "$mtls_dir/gateway.ext" >/dev/null 2>&1 + for secret_name in \ + select-ai-gateway-mtls-ca \ + select-ai-gateway-mtls-cert \ + select-ai-gateway-mtls-key \ + select-ai-worker-mtls-cert \ + select-ai-worker-mtls-key; do + gcloud secrets describe "$secret_name" --project="$project_id" >/dev/null 2>&1 || \ + gcloud secrets create "$secret_name" --replication-policy=automatic --project="$project_id" + done + gcloud secrets versions add select-ai-gateway-mtls-ca --data-file="$mtls_dir/ca.crt" --project="$project_id" + gcloud secrets versions add select-ai-gateway-mtls-cert --data-file="$mtls_dir/gateway.crt" --project="$project_id" + gcloud secrets versions add select-ai-gateway-mtls-key --data-file="$mtls_dir/gateway.key" --project="$project_id" + gcloud secrets versions add select-ai-worker-mtls-cert --data-file="$mtls_dir/worker.crt" --project="$project_id" + gcloud secrets versions add select-ai-worker-mtls-key --data-file="$mtls_dir/worker.key" --project="$project_id" + if [[ "$rotate_worker_mtls" == "true" ]]; then + echo "Rotating worker mTLS material; worker sessions will be interrupted." + else + echo "Created initial worker mTLS material." + fi + else + echo "Reusing existing worker mTLS material." + fi + if ! gcloud dns managed-zones describe "$worker_dns_zone" --project="$project_id" >/dev/null 2>&1; then + gcloud dns managed-zones create "$worker_dns_zone" --dns-name="${worker_domain}." \ + --visibility=private --networks="https://www.googleapis.com/compute/v1/projects/$project_id/global/networks/$network" \ + --description="Private DNS records for Select AI mTLS workers" \ + --project="$project_id" + fi + runtime_sa="${project_id_number:-$(gcloud projects describe "$project_id" --format='value(projectNumber)')}-compute@developer.gserviceaccount.com" + for secret_name in \ + select-ai-gateway-mtls-ca \ + select-ai-gateway-mtls-cert \ + select-ai-gateway-mtls-key \ + select-ai-worker-mtls-cert \ + select-ai-worker-mtls-key; do + gcloud secrets add-iam-policy-binding "$secret_name" --project="$project_id" \ + --member="serviceAccount:$runtime_sa" --role="roles/secretmanager.secretAccessor" >/dev/null + done +fi + +gcloud builds submit "$repo_root" \ + --project="$project_id" \ + --region="$region" \ + --config="$repo_root/gcloud/gateway/cloudbuild.yaml" \ + --substitutions="_REGION=$region,_CLUSTER=$cluster,_REPOSITORY=$repository,_IMAGE_TAG=$image_tag,_GATEWAY_SERVICE=$gateway_service,_NETWORK=$network,_SUBNET=$subnet,_WORKER_REPLICAS=$worker_replicas,_ENABLE_WORKER_MTLS=$enable_worker_mtls,_ROTATE_WORKER_MTLS=$rotate_worker_mtls,_WORKER_DOMAIN=$worker_domain,_WORKER_DNS_ZONE=$worker_dns_zone" + +gateway_url="$(gcloud run services describe "$gateway_service" \ + --region="$region" --project="$project_id" --format='value(status.url)')" +project_number="$(gcloud projects describe "$project_id" --format='value(projectNumber)')" +gemini_service_agent="service-$project_number@gcp-sa-discoveryengine.iam.gserviceaccount.com" + +gcloud run services add-iam-policy-binding "$gateway_service" \ + --region="$region" --project="$project_id" \ + --member="serviceAccount:$gemini_service_agent" \ + --role="roles/run.invoker" >/dev/null + +active_account="$(gcloud auth list --filter=status:ACTIVE --format='value(account)')" +if [[ -z "$active_account" ]]; then + echo "No active gcloud account. Run: gcloud auth login" >&2 + exit 2 +fi +if gcloud iam service-accounts describe "$active_account" \ + --project="$project_id" >/dev/null 2>&1; then + deployer_member="serviceAccount:$active_account" +else + deployer_member="user:$active_account" +fi +gcloud run services add-iam-policy-binding "$gateway_service" \ + --region="$region" --project="$project_id" \ + --member="$deployer_member" --role="roles/run.invoker" >/dev/null + +printf 'Gateway Agent Card URL:\n%s/.well-known/agent-card.json\n' "$gateway_url" diff --git a/gcloud/gateway/gke/consul.yaml b/gcloud/gateway/gke/consul.yaml new file mode 100644 index 0000000..19569c0 --- /dev/null +++ b/gcloud/gateway/gke/consul.yaml @@ -0,0 +1,50 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: consul-server + namespace: select-ai-gateway +spec: + replicas: 1 + selector: + matchLabels: + app: consul-server + template: + metadata: + labels: + app: consul-server + spec: + containers: + - name: consul + image: hashicorp/consul:1.21.5 + args: + - agent + - -server + - -bootstrap-expect=1 + - -client=0.0.0.0 + ports: + - name: http + containerPort: 8500 + readinessProbe: + httpGet: + path: /v1/status/leader + port: http + livenessProbe: + httpGet: + path: /v1/status/leader + port: http +--- +apiVersion: v1 +kind: Service +metadata: + name: consul-server + namespace: select-ai-gateway + annotations: + networking.gke.io/load-balancer-type: Internal +spec: + type: LoadBalancer + selector: + app: consul-server + ports: + - name: http + port: 8500 + targetPort: http diff --git a/gcloud/gateway/gke/namespace.yaml b/gcloud/gateway/gke/namespace.yaml new file mode 100644 index 0000000..dc1ff53 --- /dev/null +++ b/gcloud/gateway/gke/namespace.yaml @@ -0,0 +1,4 @@ +apiVersion: v1 +kind: Namespace +metadata: + name: select-ai-gateway diff --git a/gcloud/gateway/gke/worker-mtls.yaml b/gcloud/gateway/gke/worker-mtls.yaml new file mode 100644 index 0000000..9861212 --- /dev/null +++ b/gcloud/gateway/gke/worker-mtls.yaml @@ -0,0 +1,82 @@ +apiVersion: apps/v1 +kind: StatefulSet +metadata: + name: select-ai-worker + namespace: select-ai-gateway +spec: + serviceName: select-ai-worker + replicas: WORKER_REPLICAS + selector: + matchLabels: + app: select-ai-worker + template: + metadata: + labels: + app: select-ai-worker + spec: + terminationGracePeriodSeconds: 30 + containers: + - name: worker + image: WORKER_IMAGE + command: [select-ai] + args: + - a2a + - worker + - --host + - 0.0.0.0 + - --port + - "8443" + - --tls-cert-file + - /var/run/select-ai-mtls/tls.crt + - --tls-key-file + - /var/run/select-ai-mtls/tls.key + - --tls-ca-file + - /var/run/select-ai-mtls/gateway-ca.crt + ports: + - name: https + containerPort: 8443 + env: + - name: CONSUL_HTTP_URL + value: http://consul-server.select-ai-gateway.svc.cluster.local:8500 + - name: WORKER_ID + valueFrom: + fieldRef: + fieldPath: metadata.name + - name: POD_NAME + valueFrom: + fieldRef: + fieldPath: metadata.name + - name: WORKER_ADDRESS + valueFrom: + fieldRef: + fieldPath: status.podIP + - name: WORKER_PORT + value: "8443" + - name: WORKER_ENDPOINT + value: https://$(POD_NAME).WORKER_DOMAIN:8443 + volumeMounts: + - name: worker-server-tls + mountPath: /var/run/select-ai-mtls/tls.crt + subPath: tls.crt + readOnly: true + - name: worker-server-tls + mountPath: /var/run/select-ai-mtls/tls.key + subPath: tls.key + readOnly: true + - name: gateway-client-ca + mountPath: /var/run/select-ai-mtls/gateway-ca.crt + subPath: ca.crt + readOnly: true + readinessProbe: + tcpSocket: + port: https + livenessProbe: + tcpSocket: + port: https + volumes: + - name: worker-server-tls + secret: + secretName: select-ai-worker-server-tls + - name: gateway-client-ca + secret: + secretName: select-ai-gateway-client-ca diff --git a/gcloud/gateway/gke/worker.yaml b/gcloud/gateway/gke/worker.yaml new file mode 100644 index 0000000..fc24c5e --- /dev/null +++ b/gcloud/gateway/gke/worker.yaml @@ -0,0 +1,52 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: select-ai-worker + namespace: select-ai-gateway +spec: + replicas: WORKER_REPLICAS + selector: + matchLabels: + app: select-ai-worker + template: + metadata: + labels: + app: select-ai-worker + spec: + terminationGracePeriodSeconds: 30 + containers: + - name: worker + image: WORKER_IMAGE + command: + - select-ai + args: + - a2a + - worker + - --host + - 0.0.0.0 + - --port + - "8080" + ports: + - name: http + containerPort: 8080 + env: + - name: CONSUL_HTTP_URL + value: http://consul-server.select-ai-gateway.svc.cluster.local:8500 + - name: WORKER_ID + valueFrom: + fieldRef: + fieldPath: metadata.name + - name: WORKER_ADDRESS + valueFrom: + fieldRef: + fieldPath: status.podIP + - name: WORKER_PORT + value: "8080" + readinessProbe: + httpGet: + path: /health + port: http + livenessProbe: + httpGet: + path: /health + port: http diff --git a/gcloud/standalone/README.md b/gcloud/standalone/README.md new file mode 100644 index 0000000..5c2df66 --- /dev/null +++ b/gcloud/standalone/README.md @@ -0,0 +1,159 @@ +# Standalone Select AI A2A deployment + +`gcloud/standalone/deploy.sh` builds or selects a Select AI container image, creates or +updates a private Cloud Run service, and configures its database secrets. Run +it on a machine with the [Google Cloud CLI](https://cloud.google.com/sdk/docs/install) +installed and authenticated to the target project. + +## IAM permissions + +The scripts use the active `gcloud` identity. They do not elevate its access. + +### Deployer (the active gcloud identity) + +| Operation | Required permissions | +| --- | --- | +| Inspect and create the Artifact Registry repository | `artifactregistry.repositories.get`, `artifactregistry.repositories.create` | +| Inspect and create the default runtime service account | `iam.serviceAccounts.get`, `iam.serviceAccounts.create` | +| Deploy or update Cloud Run | `run.services.create`, `run.services.update`, `run.services.get`, `run.operations.get`; `iam.serviceAccounts.actAs` on the runtime service account; `artifactregistry.repositories.downloadArtifacts` on the image repository | +| With `--build`, upload local source, submit, and wait for a build | `storage.buckets.get`, `storage.objects.create` on the configured source-staging bucket; `cloudbuild.builds.create`, `cloudbuild.builds.get`, `serviceusage.services.use` | +| Inspect, create, and add versions to database or wallet secrets | `secretmanager.secrets.get`, `secretmanager.secrets.create`, `secretmanager.versions.add` | +| Grant the runtime account access to those secrets | `secretmanager.secrets.getIamPolicy`, `secretmanager.secrets.setIamPolicy` | +| Grant Gemini Enterprise and the active gcloud identity access to the service | `run.services.getIamPolicy`, `run.services.setIamPolicy` | +| Obtain the project number | `resourcemanager.projects.get` | + +### Runtime service account + +| Operation | Required permissions | +| --- | --- | +| Read database and wallet secrets while serving requests | `secretmanager.versions.access` | + +### Other service identities + +| Principal | Operation | Required permissions | +| --- | --- | --- | +| Cloud Build execution service account | With `--build`, push the built image | `artifactregistry.repositories.uploadArtifacts` | +| Gemini Enterprise service agent | Invoke the private Cloud Run service | `run.routes.invoke` | +| Active gcloud identity | Fetch the Agent Card after deployment | `run.routes.invoke` | + +The source-staging bucket is Cloud Build's default unless a custom bucket is +configured. Cloud Build also needs access to its build-log destination; the +default same-project build account has that access. If your organization uses +a custom build service account, source bucket, or log bucket, its administrator +must grant the equivalent Cloud Storage permissions on those resources. + +Google Cloud references: [Service Usage access control](https://cloud.google.com/service-usage/docs/access-control), [Cloud Run deployment permissions](https://cloud.google.com/run/docs/reference/iam/roles), [Secret Manager access control](https://cloud.google.com/secret-manager/docs/access-control), [Artifact Registry roles](https://cloud.google.com/artifact-registry/docs/access-control), and [Cloud Build roles](https://cloud.google.com/build/docs/iam-roles-permissions). + +The rows that set IAM policy are administrative mutations. They are present +because `deploy.sh` creates and rotates secrets and configures private-service +invocation. If your customer deployment identity must not change IAM, provision +the secrets and the `secretmanager.versions.access`/`run.routes.invoke` +permissions beforehand, then +remove those policy-setting commands from the deployment workflow. + +## Prerequisite: enable project APIs once + +An administrator must enable these APIs once for the project: + +```bash +gcloud services enable \ + run.googleapis.com \ + cloudbuild.googleapis.com \ + artifactregistry.googleapis.com \ + secretmanager.googleapis.com \ + discoveryengine.googleapis.com \ + --project PROJECT_ID +``` + +## Deploy (and update) the A2A server + +```bash +gcloud/standalone/deploy.sh --build +``` + +On the first deployment, the script prompts for the ADB user, password, and +connect descriptor. It stores them in Secret Manager under names based on the +Cloud Run service, and grants only the runtime service account access. The +container receives the values as `SELECT_AI_USER`, `SELECT_AI_PASSWORD`, and +`SELECT_AI_DB_CONNECT_STRING`; they are never placed in the image or source +tree. + +### Optional: Autonomous Database mTLS wallet + +The Select AI SDK already supports `wallet_location` and `wallet_password`. +For Cloud Run, pass the path to the downloaded Autonomous Database wallet ZIP +on the first deployment (or when replacing it): + +```bash +gcloud/standalone/deploy.sh --wallet-archive /path/to/Wallet_database.zip +``` + +The script prompts for the wallet password, stores the ZIP and password as +service-specific Secret Manager secrets, and grants access only to the runtime +service account. Cloud Run mounts the ZIP read-only; its A2A launcher expands it +into ephemeral `/tmp` storage before starting the SDK, verifies it contains +`ewallet.pem`, and sets `SELECT_AI_WALLET_LOCATION` to that file's directory. +Do not commit the wallet ZIP or put its contents in the image. + +Later deploys reuse the wallet. To replace it, pass `--wallet-archive` again. + +The first deployment needs `--build` (or an explicit `--image-uri`). Later +deployments reuse the image already deployed to the service, so changing Cloud +Run configuration or secrets does not create another image. The command +deploys private Cloud Run, sets the final public URL in the Agent Card, grants your active +gcloud identity and Gemini Enterprise Discovery Engine service agent the +`run.routes.invoke` permission for this Cloud Run service. + +The default Cloud Run service is `oracle-a2a-agent`. Its default Agent Team, +installed in Oracle Database, is `ORACLE_AI_DATABASE_AGENT`. Override either +with explicit options: + +```bash +gcloud/standalone/deploy.sh --service sales-analyst-a2a --a2a-team SALES_ANALYST +``` + +Use a distinct `--service` value for each A2A team. Each service gets distinct Secret +Manager secret names by default, so credentials remain attached to that A2A +server. + +`--max-instances` controls the number of Cloud Run containers. Each container +can use up to 10 Oracle connections by default; change that limit with +`--pool-max-size`, for example `gcloud/standalone/deploy.sh --pool-max-size 20`. + +### Update the Select AI SDK or this repository + +Update the checkout (or modify its dependency version), then explicitly build +and deploy the new image: + +```bash +git pull +gcloud/standalone/deploy.sh --build +``` + +`--build` creates a freshly tagged image from the current source; without it, +the existing image is reused. Existing database secrets are reused without +prompting. To rotate the ADB credentials, explicitly request it: + +```bash +gcloud/standalone/deploy.sh --rotate-db-credentials +``` + +### What `cloudbuild.yaml` does + +`gcloud/standalone/deploy.sh --build` uses `gcloud/standalone/cloudbuild.yaml` to tell Cloud Build to build +`docker/Dockerfile` and push it to Artifact Registry. It is build configuration, +not a command you run. The build context is the repository root, so the image +can install the Select AI source from `pyproject.toml` and `src/`. + +### Cloud Build upload contents + +Before the build starts, `gcloud builds submit` archives and uploads the +repository root. The root `.gcloudignore` excludes local virtual environments, +generated documentation, test data, caches, credentials, and Git metadata. +Keep `src/`, `pyproject.toml`, `docker/`, and `gcloud/` in the upload; they are +required to build the image. If the upload is unexpectedly large, check local +directories against `.gcloudignore` before running `--build` again. + +After a successful deployment, the script prints the A2A Agent Card JSON. +Paste that JSON into Gemini Enterprise to register the private service. The +required Gemini Enterprise invocation permission has already been added. diff --git a/gcloud/cloudbuild.yaml b/gcloud/standalone/cloudbuild.yaml similarity index 82% rename from gcloud/cloudbuild.yaml rename to gcloud/standalone/cloudbuild.yaml index 9977f13..3cedbda 100644 --- a/gcloud/cloudbuild.yaml +++ b/gcloud/standalone/cloudbuild.yaml @@ -5,8 +5,8 @@ # http://oss.oracle.com/licenses/upl. # ----------------------------------------------------------------------------- -# Build one reusable A2A server image. Database team selection is Cloud Run -# configuration, not an image-build input. +# Build one reusable Select AI image for standalone servers, the gateway, and +# GKE worker replicas. Runtime role selection is deployment configuration. steps: - name: gcr.io/cloud-builders/docker args: diff --git a/gcloud/deploy.sh b/gcloud/standalone/deploy.sh similarity index 98% rename from gcloud/deploy.sh rename to gcloud/standalone/deploy.sh index 0eeab66..3ac7ff1 100755 --- a/gcloud/deploy.sh +++ b/gcloud/standalone/deploy.sh @@ -15,7 +15,7 @@ set -euo pipefail usage() { cat <<'EOF' -Usage: gcloud/deploy.sh [options] +Usage: gcloud/standalone/deploy.sh [options] Deploy the Select AI A2A server to private Cloud Run. @@ -45,7 +45,7 @@ Options: EOF } -repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" project_id="" region="us-central1" repository="select-ai" @@ -227,7 +227,7 @@ if [[ "$build_image" == true ]]; then image_uri="$region-docker.pkg.dev/$project_id/$repository/select-ai:$image_tag" echo "Building $image_uri" gcloud builds submit "$repo_root" --project="$project_id" \ - --config="$repo_root/gcloud/cloudbuild.yaml" \ + --config="$repo_root/gcloud/standalone/cloudbuild.yaml" \ --substitutions="_REGION=$region,_REPOSITORY=$repository,_IMAGE_TAG=$image_tag" elif [[ -z "$image_uri" && "$service_exists" == true ]]; then image_uri="$(gcloud run services describe "$service" --project="$project_id" --region="$region" \ diff --git a/pyproject.toml b/pyproject.toml index 1c39728..6ffa295 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,6 +48,9 @@ cli = [ "click", "a2a-sdk[http-server]>=1.0.3", "uvicorn[standard]>=0.30", + "fastapi>=0.115,<1", + "httpx>=0.27,<1", + "requests>=2.32,<3", ] a2a = [ "select_ai[cli]", diff --git a/src/select_ai/agent/a2a/__init__.py b/src/select_ai/agent/a2a/__init__.py index f4786b6..7424546 100644 --- a/src/select_ai/agent/a2a/__init__.py +++ b/src/select_ai/agent/a2a/__init__.py @@ -5,4 +5,32 @@ # https://oss.oracle.com/licenses/upl. # ----------------------------------------------------------------------------- -"""A2A support for Select AI Agent Teams.""" +"""A2A support for Select AI Agent Teams and temporary sessions.""" + +from .models import GatewaySettings, SessionInfo, SessionRoute + + +def create_gateway_app(settings): + """Build the public A2A/A2UI gateway application.""" + from select_ai.agent.a2a.gateway import create_gateway_app as factory + + return factory(settings) + + +def create_worker_app( + session_ttl_seconds: int = 900, + session_start_timeout_seconds: int = 30, +): + """Build the internal worker application.""" + from select_ai.agent.a2a.worker import create_worker_app as factory + + return factory(session_ttl_seconds, session_start_timeout_seconds) + + +__all__ = [ + "GatewaySettings", + "SessionInfo", + "SessionRoute", + "create_gateway_app", + "create_worker_app", +] diff --git a/src/select_ai/agent/a2a/a2ui.py b/src/select_ai/agent/a2a/a2ui.py new file mode 100644 index 0000000..a25b7f9 --- /dev/null +++ b/src/select_ai/agent/a2a/a2ui.py @@ -0,0 +1,33 @@ +# ----------------------------------------------------------------------------- +# Copyright (c) 2026, Oracle and/or its affiliates. +# +# Licensed under the Universal Permissive License v 1.0 as shown at +# https://oss.oracle.com/licenses/upl. +# ----------------------------------------------------------------------------- + +"""Shared A2UI protocol declarations for Select AI A2A agents.""" + +from a2a.types import AgentExtension +from google.protobuf.json_format import ParseDict +from google.protobuf.struct_pb2 import Struct + +from select_ai.agent.a2a.forms import _CATALOG + +A2UI_EXTENSION_URI = "https://a2ui.org/a2a-extension/a2ui/v0.9" +A2UI_MIME_TYPE = "application/json+a2ui" + + +def a2ui_extension() -> AgentExtension: + """Return the A2UI v0.9 capability used by Gemini Enterprise.""" + params = ParseDict( + { + "acceptsInlineCatalogs": True, + "supportedCatalogIds": [_CATALOG], + }, + Struct(), + ) + return AgentExtension( + uri=A2UI_EXTENSION_URI, + description="Provides agent driven UI using the A2UI JSON format.", + params=params, + ) diff --git a/src/select_ai/agent/a2a/forms.py b/src/select_ai/agent/a2a/forms.py new file mode 100644 index 0000000..6823464 --- /dev/null +++ b/src/select_ai/agent/a2a/forms.py @@ -0,0 +1,117 @@ +# ----------------------------------------------------------------------------- +# Copyright (c) 2026, Oracle and/or its affiliates. +# +# Licensed under the Universal Permissive License v 1.0 as shown at +# https://oss.oracle.com/licenses/upl. +# ----------------------------------------------------------------------------- + +"""A2UI connection form emitted by the public gateway.""" + + +_CATALOG = ( + "https://www.gstatic.com/vertexaisearch/a2ui/v0_9/" + "gemini_enterprise_composite_catalog.json" +) + + +def connection_form() -> list[dict]: + """Return the non-persistent database connection form.""" + return [ + { + "version": "v0.9", + "createSurface": { + "surfaceId": "db-connect", + "catalogId": _CATALOG, + }, + }, + { + "version": "v0.9", + "updateComponents": { + "surfaceId": "db-connect", + "components": [ + {"id": "root", "component": "Card", "child": "column"}, + { + "id": "column", + "component": "Column", + "children": [ + "title", + "dsn", + "user", + "password", + "team", + "connect", + ], + }, + { + "id": "title", + "component": "Text", + "text": "Connect to Oracle Database", + "variant": "h2", + }, + { + "id": "dsn", + "component": "TextField", + "label": "Database DSN", + "value": {"path": "/dsn"}, + "variant": "shortText", + }, + { + "id": "user", + "component": "TextField", + "label": "Database username", + "value": {"path": "/username"}, + "variant": "shortText", + }, + { + "id": "password", + "component": "TextField", + "label": "Database password", + "value": {"path": "/password"}, + "variant": "obscured", + }, + { + "id": "team", + "component": "TextField", + "label": "Select AI team name", + "value": {"path": "/team_name"}, + "variant": "shortText", + }, + { + "id": "connect_label", + "component": "Text", + "text": "Connect", + }, + { + "id": "connect", + "component": "Button", + "child": "connect_label", + "variant": "primary", + "action": { + "event": { + "name": "submit_database_connection", + "context": { + "dsn": {"path": "/dsn"}, + "username": {"path": "/username"}, + "password": {"path": "/password"}, + "team_name": {"path": "/team_name"}, + }, + } + }, + }, + ], + }, + }, + { + "version": "v0.9", + "updateDataModel": { + "surfaceId": "db-connect", + "path": "/", + "value": { + "dsn": "", + "username": "", + "password": "", + "team_name": "", + }, + }, + }, + ] diff --git a/src/select_ai/agent/a2a/gateway.py b/src/select_ai/agent/a2a/gateway.py new file mode 100644 index 0000000..4b0957e --- /dev/null +++ b/src/select_ai/agent/a2a/gateway.py @@ -0,0 +1,231 @@ +# ----------------------------------------------------------------------------- +# Copyright (c) 2026, Oracle and/or its affiliates. +# +# Licensed under the Universal Permissive License v 1.0 as shown at +# https://oss.oracle.com/licenses/upl. +# ----------------------------------------------------------------------------- + +"""Public A2A/A2UI gateway for temporary Select AI database sessions.""" + +from __future__ import annotations + +import asyncio + +from a2a.compat.v0_3.conversions import to_compat_agent_card +from a2a.helpers import ( + new_data_part, + new_task_from_user_message, + new_text_part, +) +from a2a.server.agent_execution import AgentExecutor +from a2a.server.request_handlers import DefaultRequestHandler +from a2a.server.routes import create_jsonrpc_routes +from a2a.server.tasks import InMemoryTaskStore, TaskUpdater +from a2a.types import ( + AgentCapabilities, + AgentCard, + AgentInterface, + AgentSkill, +) +from google.protobuf.json_format import MessageToDict, ParseDict +from starlette.applications import Starlette +from starlette.responses import JSONResponse +from starlette.routing import Route + +from select_ai.agent.a2a.a2ui import ( + A2UI_EXTENSION_URI, + A2UI_MIME_TYPE, + a2ui_extension, +) +from select_ai.agent.a2a.forms import connection_form +from select_ai.agent.a2a.models import GatewaySettings, SessionInfo +from select_ai.agent.a2a.results import add_team_result +from select_ai.agent.a2a.worker_client import ReconnectRequired, WorkerClient +from select_ai.version import __version__ + + +class GatewayExecutor(AgentExecutor): + """Route each A2A context to one in-memory worker session.""" + + def __init__(self, worker_client: WorkerClient): + self.worker_client = worker_client + self.sessions: dict[str, str] = {} + + async def execute(self, context, event_queue): + task = context.current_task or new_task_from_user_message( + context.message + ) + if context.current_task is None: + await event_queue.enqueue_event(task) + context_id = task.context_id or task.id + updater = TaskUpdater( + event_queue=event_queue, + task_id=task.id, + context_id=task.context_id, + ) + await updater.start_work() + action = self._a2ui_action(context) + if action and action.get("name") == "submit_database_connection": + parts = await self._open_session( + action.get("context", {}), + context_id, + ) + artifact_name = "database-session" + extensions = None + elif context_id not in self.sessions: + parts = [_a2ui_part(message) for message in connection_form()] + artifact_name = "database-connection-form" + extensions = [A2UI_EXTENSION_URI] + else: + result = await self._send_prompt( + context_id, + context.get_user_input(), + ) + await add_team_result(updater, result) + await updater.complete() + return + await updater.add_artifact( + parts=parts, + name=artifact_name, + last_chunk=True, + extensions=extensions, + ) + await updater.complete() + + async def _open_session(self, action_context: dict, context_id: str): + try: + session_info = SessionInfo.from_a2ui_event(action_context) + session_id = await asyncio.to_thread( + self.worker_client.open_session, + context_id, + session_info, + ) + except (ValueError, RuntimeError): + return [ + new_text_part( + "Could not connect. Check the DSN, credentials, and team " + "name." + ) + ] + self.sessions[context_id] = session_id + return [new_text_part("Connected. Ask a database question.")] + + async def _send_prompt(self, context_id: str, prompt: str) -> str | None: + try: + result = await asyncio.to_thread( + self.worker_client.send_prompt, + self.sessions[context_id], + prompt, + ) + except ReconnectRequired: + self.sessions.pop(context_id, None) + return "Your database session ended. Please reconnect." + return result + + async def cancel(self, context, event_queue): + """Close the live worker session when the A2A task is cancelled.""" + task = context.current_task + if task is None: + return + context_id = task.context_id or task.id + session_id = self.sessions.pop(context_id, None) + if session_id: + await asyncio.to_thread( + self.worker_client.close_session, + session_id, + ) + updater = TaskUpdater( + event_queue=event_queue, + task_id=task.id, + context_id=task.context_id, + ) + await updater.cancel() + + @staticmethod + def _a2ui_action(context) -> dict | None: + for part in context.message.parts: + if part.WhichOneof("content") != "data": + continue + data = MessageToDict(part.data) + messages = data if isinstance(data, list) else [data] + for message in messages: + if ( + not isinstance(message, dict) + or message.get("version") != "v0.9" + ): + continue + action = message.get("action") + if isinstance(action, dict): + return action + return None + + +def _a2ui_part(message: dict): + """Encode one A2UI operation in the format used by Gemini Enterprise.""" + part = new_data_part(message) + ParseDict({"mimeType": A2UI_MIME_TYPE}, part.metadata) + return part + + +def create_gateway_app(settings: GatewaySettings) -> Starlette: + """Build a Gemini Enterprise-compatible A2A v0.3 gateway application.""" + description = "Connects a user to a temporary Select AI database session." + endpoint = f"{settings.agent_url}/a2a/jsonrpc/" + card = AgentCard( + name="Select AI Database Gateway", + description=description, + version=__version__, + default_input_modes=["text/plain", A2UI_MIME_TYPE], + default_output_modes=["text/plain", A2UI_MIME_TYPE], + capabilities=AgentCapabilities( + streaming=False, + extensions=[a2ui_extension()], + ), + supported_interfaces=[ + AgentInterface( + protocol_binding="JSONRPC", + protocol_version="1.0", + url=endpoint, + ), + AgentInterface( + protocol_binding="JSONRPC", + protocol_version="0.3", + url=endpoint, + ), + ], + skills=[ + AgentSkill( + id="database_connect", + name="Connect to database", + description=description, + tags=["oracle", "select-ai"], + examples=[], + input_modes=["text/plain", A2UI_MIME_TYPE], + output_modes=["text/plain", A2UI_MIME_TYPE], + ) + ], + ) + handler = DefaultRequestHandler( + agent_executor=GatewayExecutor(WorkerClient(settings)), + task_store=InMemoryTaskStore(), + agent_card=card, + ) + compat_card = to_compat_agent_card(card).model_dump( + by_alias=True, + exclude_none=True, + ) + + async def get_agent_card(_request): + return JSONResponse(compat_card) + + routes = [ + Route("/.well-known/agent-card.json", get_agent_card, methods=["GET"]) + ] + routes.extend( + create_jsonrpc_routes( + handler, + rpc_url="/a2a/jsonrpc/", + enable_v0_3_compat=True, + ) + ) + return Starlette(routes=routes) diff --git a/src/select_ai/agent/a2a/models.py b/src/select_ai/agent/a2a/models.py new file mode 100644 index 0000000..2812f61 --- /dev/null +++ b/src/select_ai/agent/a2a/models.py @@ -0,0 +1,71 @@ +# ----------------------------------------------------------------------------- +# Copyright (c) 2026, Oracle and/or its affiliates. +# +# Licensed under the Universal Permissive License v 1.0 as shown at +# https://oss.oracle.com/licenses/upl. +# ----------------------------------------------------------------------------- + +"""Configuration and transient request models for the A2A runtime.""" + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class GatewaySettings: + """Configuration for the public gateway process.""" + + agent_url: str + consul_url: str + worker_service: str + session_ttl_seconds: int + worker_tls_ca_file: str | None = None + worker_tls_cert_file: str | None = None + worker_tls_key_file: str | None = None + + def __post_init__(self) -> None: + if self.session_ttl_seconds < 1: + raise ValueError("session_ttl_seconds must be at least 1") + tls_files = ( + self.worker_tls_ca_file, + self.worker_tls_cert_file, + self.worker_tls_key_file, + ) + if any(tls_files) and not all(tls_files): + raise ValueError( + "worker mTLS requires a CA file, client certificate, and " + "client key." + ) + object.__setattr__(self, "agent_url", self.agent_url.rstrip("/")) + object.__setattr__(self, "consul_url", self.consul_url.rstrip("/")) + + @property + def worker_mtls_enabled(self) -> bool: + """Whether gateway-to-worker calls require mutual TLS.""" + return self.worker_tls_ca_file is not None + + +@dataclass(frozen=True) +class SessionInfo: + """Credentials used only while opening one in-memory worker session.""" + + dsn: str + username: str + password: str + team_name: str + + @classmethod + def from_a2ui_event(cls, event: dict) -> "SessionInfo": + required = ("dsn", "username", "password", "team_name") + if not all( + isinstance(event.get(key), str) and event[key] for key in required + ): + raise ValueError("All database connection fields are required.") + return cls(**{key: event[key] for key in required}) + + +@dataclass(frozen=True) +class SessionRoute: + """Non-secret Consul record that routes a session to one worker.""" + + endpoint: str + expires_at: float diff --git a/src/select_ai/agent/a2a/results.py b/src/select_ai/agent/a2a/results.py new file mode 100644 index 0000000..7e5c0a4 --- /dev/null +++ b/src/select_ai/agent/a2a/results.py @@ -0,0 +1,58 @@ +# ----------------------------------------------------------------------------- +# Copyright (c) 2026, Oracle and/or its affiliates. +# +# Licensed under the Universal Permissive License v 1.0 as shown at +# https://oss.oracle.com/licenses/upl. +# ----------------------------------------------------------------------------- + +"""Translate serialized Select AI team results into A2A response artifacts.""" + +from __future__ import annotations + +import json + +from a2a.helpers import new_data_part, new_text_part +from a2a.server.tasks import TaskUpdater +from google.protobuf.json_format import ParseDict + + +def message_parts(result: str | None): + """Convert a serialized A2A message into its constituent parts.""" + if not result: + return None + try: + payload = json.loads(result) + except json.JSONDecodeError: + return None + if not isinstance(payload, dict): + return None + serialized_parts = payload.get("parts") + if payload.get("kind") != "message" or not isinstance( + serialized_parts, + list, + ): + return None + parts = [] + for part in serialized_parts: + if not isinstance(part, dict): + return None + if part.get("kind") == "text" and isinstance(part.get("text"), str): + parts.append(new_text_part(part["text"])) + elif part.get("kind") == "data" and "data" in part: + output_part = new_data_part(part["data"]) + if isinstance(part.get("metadata"), dict): + ParseDict(part["metadata"], output_part.metadata) + parts.append(output_part) + else: + # Avoid silently discarding an unsupported part type. + return None + return parts or None + + +async def add_team_result(updater: TaskUpdater, result: str | None) -> None: + """Attach a raw team result as the completed database-result artifact.""" + await updater.add_artifact( + parts=message_parts(result) or [new_text_part(result or "")], + name="database-agent-result", + last_chunk=True, + ) diff --git a/src/select_ai/agent/a2a/server.py b/src/select_ai/agent/a2a/server.py index 6b28a26..d2c06fa 100644 --- a/src/select_ai/agent/a2a/server.py +++ b/src/select_ai/agent/a2a/server.py @@ -7,16 +7,11 @@ """A2A HTTP server for Oracle Database AI Agent Teams.""" -import json from contextlib import asynccontextmanager from typing import Optional from a2a.compat.v0_3.conversions import to_compat_agent_card -from a2a.helpers import ( - new_data_part, - new_task_from_user_message, - new_text_part, -) +from a2a.helpers import new_task_from_user_message from a2a.server.agent_execution import AgentExecutor from a2a.server.request_handlers import DefaultRequestHandler from a2a.server.routes import create_jsonrpc_routes @@ -28,29 +23,12 @@ import select_ai from select_ai.agent import AsyncTeam +from select_ai.agent.a2a.a2ui import A2UI_MIME_TYPE, a2ui_extension from select_ai.agent.a2a.context_store import OracleContextStore +from select_ai.agent.a2a.results import add_team_result from select_ai.agent.a2a.task_store import OracleTaskStore from select_ai.version import __version__ -_A2UI_MIME_TYPE = "application/a2ui+json" - - -def _a2ui_payload(result: str | None) -> dict | None: - """Return an A2UI response envelope, if ``RUN_TEAM`` returned one.""" - if not result: - return None - try: - payload = json.loads(result) - except json.JSONDecodeError: - return None - if not isinstance(payload, dict): - return None - if payload.get("metadata", {}).get("mimeType") != _A2UI_MIME_TYPE: - return None - if not isinstance(payload.get("data"), list): - return None - return payload - class DatabaseTeamExecutor(AgentExecutor): """Execute A2A requests with one Oracle conversation per A2A context.""" @@ -81,16 +59,7 @@ async def execute(self, context, event_queue): prompt=context.get_user_input(), params={"conversation_id": conversation_id}, ) - a2ui_payload = _a2ui_payload(result) - await updater.add_artifact( - parts=( - [new_data_part(a2ui_payload)] - if a2ui_payload is not None - else [new_text_part(result or "")] - ), - name="database-agent-result", - last_chunk=True, - ) + await add_team_result(updater, result) await updater.complete() async def cancel(self, context, event_queue): @@ -184,8 +153,11 @@ def _build_agent_card( description=description, version=__version__, default_input_modes=["text/plain"], - default_output_modes=["text/plain"], - capabilities=AgentCapabilities(streaming=True), + default_output_modes=["text/plain", A2UI_MIME_TYPE], + capabilities=AgentCapabilities( + streaming=True, + extensions=[a2ui_extension()], + ), supported_interfaces=[ AgentInterface( protocol_binding="JSONRPC", @@ -206,7 +178,7 @@ def _build_agent_card( tags=["oracle", "database", "select-ai"], examples=[], input_modes=["text/plain"], - output_modes=["text/plain"], + output_modes=["text/plain", A2UI_MIME_TYPE], ) ], ) diff --git a/src/select_ai/agent/a2a/worker.py b/src/select_ai/agent/a2a/worker.py new file mode 100644 index 0000000..d6e7b04 --- /dev/null +++ b/src/select_ai/agent/a2a/worker.py @@ -0,0 +1,436 @@ +# ----------------------------------------------------------------------------- +# Copyright (c) 2026, Oracle and/or its affiliates. +# +# Licensed under the Universal Permissive License v 1.0 as shown at +# https://oss.oracle.com/licenses/upl. +# ----------------------------------------------------------------------------- + +"""Internal worker that owns temporary Select AI child processes.""" + +from __future__ import annotations + +import asyncio +import contextlib +import logging +import multiprocessing +import os +import socket +import time +from contextlib import asynccontextmanager +from dataclasses import dataclass +from multiprocessing.connection import Connection + +import httpx +from fastapi import FastAPI, HTTPException, status +from fastapi.responses import PlainTextResponse +from pydantic import BaseModel, Field, SecretStr + +LOGGER = logging.getLogger(__name__) + + +class OpenSessionRequest(BaseModel): + """Sensitive request accepted only on the internal worker interface.""" + + session_id: str + dsn: str = Field(min_length=1, max_length=4_000) + username: str = Field(min_length=1, max_length=128) + password: SecretStr = Field(min_length=1, max_length=1_024) + team_name: str = Field(min_length=1, max_length=128) + + +class PromptRequest(BaseModel): + """One user prompt for a previously opened session.""" + + prompt: str = Field(min_length=1, max_length=32_000) + + +@dataclass +class ChildSession: + """In-memory ownership record for one database-session process.""" + + process: multiprocessing.Process + connection: Connection + expires_at: float + lock: asyncio.Lock + + +class SessionWorker: + """Own one isolated Select AI runtime process for each session.""" + + def __init__( + self, + session_ttl_seconds: int, + session_start_timeout_seconds: int, + ) -> None: + self.session_ttl_seconds = session_ttl_seconds + self.session_start_timeout_seconds = session_start_timeout_seconds + self.sessions: dict[str, ChildSession] = {} + self.lock = asyncio.Lock() + + async def open(self, request: OpenSessionRequest) -> None: + """Start a child runtime and wait until its database pool is ready.""" + parent_connection, child_connection = multiprocessing.Pipe() + credentials = { + "user": request.username, + "password": request.password.get_secret_value(), + "dsn": request.dsn, + } + process = multiprocessing.Process( + target=_session_process_main, + args=( + child_connection, + credentials, + request.session_id, + request.team_name, + ), + daemon=True, + ) + process.start() + child_connection.close() + session = ChildSession( + process=process, + connection=parent_connection, + expires_at=time.monotonic() + self.session_ttl_seconds, + lock=asyncio.Lock(), + ) + try: + await self._wait_ready(session, request) + except Exception: + await self._terminate(session) + raise + async with self.lock: + previous = self.sessions.pop(request.session_id, None) + if previous: + await self._terminate(previous) + self.sessions[request.session_id] = session + + async def get(self, session_id: str) -> ChildSession: + """Return a live session, closing it if its expiry has elapsed.""" + async with self.lock: + session = self.sessions.get(session_id) + if session and ( + session.expires_at <= time.monotonic() + or not session.process.is_alive() + ): + self.sessions.pop(session_id, None) + await self._terminate(session) + session = None + if session is None: + raise HTTPException( + status_code=404, + detail="Database session expired; reconnect required.", + ) + return session + + async def send_prompt(self, session_id: str, prompt: str) -> str | None: + """Run one prompt in the process that owns this database session.""" + session = await self.get(session_id) + async with session.lock: + if not session.process.is_alive(): + await self._discard(session_id, session) + raise HTTPException( + status_code=404, + detail="Database session expired; reconnect required.", + ) + try: + await asyncio.to_thread( + session.connection.send, + {"type": "run", "prompt": prompt}, + ) + response = await self._receive(session, timeout_seconds=120) + except (EOFError, OSError, TimeoutError) as error: + await self._discard(session_id, session) + raise HTTPException( + status_code=502, + detail=( + "Database session is unavailable; reconnect required." + ), + ) from error + if response.get("type") == "result": + return response.get("result") + LOGGER.error("Select AI session process reported a command failure.") + raise HTTPException( + status_code=502, + detail="Database session is unavailable; reconnect required.", + ) + + async def close(self, session_id: str) -> None: + """Terminate a session on an explicit gateway DELETE request.""" + async with self.lock: + session = self.sessions.pop(session_id, None) + if session is None: + raise HTTPException( + status_code=404, + detail="Database session expired; reconnect required.", + ) + await self._terminate(session) + + async def close_all(self) -> None: + """Terminate all child sessions during worker shutdown.""" + async with self.lock: + sessions = list(self.sessions.values()) + self.sessions.clear() + for session in sessions: + await self._terminate(session) + + async def _wait_ready( + self, + session: ChildSession, + request: OpenSessionRequest, + ) -> None: + try: + response = await self._receive( + session, + timeout_seconds=self.session_start_timeout_seconds, + ) + except (EOFError, OSError, TimeoutError) as error: + raise HTTPException( + status_code=504, + detail="Database session start timed out.", + ) from error + if response.get("type") == "ready": + return + detail = response.get("detail", "Database login failed.") + for secret in ( + request.password.get_secret_value(), + request.username, + request.dsn, + ): + detail = detail.replace(secret, "[REDACTED]") + LOGGER.error("Select AI session startup failed: %s", detail[-2_000:]) + raise HTTPException(status_code=400, detail="Database login failed.") + + @staticmethod + async def _receive( + session: ChildSession, + timeout_seconds: float, + ) -> dict: + available = await asyncio.to_thread( + session.connection.poll, + timeout_seconds, + ) + if not available: + raise TimeoutError() + return await asyncio.to_thread(session.connection.recv) + + async def _discard(self, session_id: str, session: ChildSession) -> None: + async with self.lock: + if self.sessions.get(session_id) is session: + self.sessions.pop(session_id, None) + await self._terminate(session) + + @staticmethod + async def _terminate(session: ChildSession) -> None: + def stop() -> None: + try: + if session.process.is_alive(): + with contextlib.suppress(OSError): + session.connection.send({"type": "close"}) + session.process.join(timeout=5) + if session.process.is_alive(): + session.process.terminate() + session.process.join(timeout=5) + finally: + session.connection.close() + + await asyncio.to_thread(stop) + + +def _session_process_main( + connection: Connection, + credentials: dict[str, str], + session_id: str, + team_name: str, +) -> None: + """Entrypoint for a child that owns one Select AI database session.""" + try: + asyncio.run( + _run_session_process( + connection, + credentials, + session_id, + team_name, + ) + ) + finally: + connection.close() + + +async def _run_session_process( + connection: Connection, + credentials: dict[str, str], + session_id: str, + team_name: str, +) -> None: + """Open one async connection and execute raw ``AsyncTeam`` prompts.""" + import select_ai + from select_ai.agent import AsyncTeam + + ready = False + try: + await select_ai.async_connect( + user=credentials["user"], + password=credentials["password"], + dsn=credentials["dsn"], + ) + if not await select_ai.async_is_connected(): + raise RuntimeError("Database login failed.") + connection.send({"type": "ready"}) + ready = True + conversation_id = None + while True: + try: + command = await asyncio.to_thread(connection.recv) + except EOFError: + return + if command.get("type") == "close": + return + if command.get("type") != "run": + connection.send( + {"type": "error", "detail": "Invalid command."} + ) + continue + try: + if conversation_id is None: + conversation = select_ai.AsyncConversation( + attributes=select_ai.ConversationAttributes( + title=f"A2A {team_name}", + description=f"Temporary session {session_id}", + ) + ) + conversation_id = await conversation.create() + result = await AsyncTeam(team_name=team_name).run( + prompt=command["prompt"], + params={"conversation_id": conversation_id}, + ) + connection.send({"type": "result", "result": result}) + except Exception: + # Keep session-process failures private from gateway callers. + LOGGER.error("Select AI session command failed") + connection.send({"type": "error"}) + except Exception as error: + LOGGER.error("Select AI session process startup failed") + if not ready: + with contextlib.suppress(OSError): + connection.send({"type": "error", "detail": str(error)}) + finally: + with contextlib.suppress(Exception): + await select_ai.async_disconnect() + + +async def _register_with_consul( + consul_url: str, + worker_id: str, + worker_address: str, + worker_port: int, + worker_endpoint: str | None = None, +) -> None: + payload = { + "Name": "select-ai-worker", + "ID": worker_id, + "Address": worker_address, + "Port": worker_port, + "Check": {"TTL": "30s", "DeregisterCriticalServiceAfter": "1m"}, + } + if worker_endpoint: + payload["Meta"] = {"endpoint": worker_endpoint.rstrip("/")} + deadline = time.monotonic() + 30 + while True: + try: + async with httpx.AsyncClient(timeout=10) as client: + response = await client.put( + f"{consul_url}/v1/agent/service/register", + json=payload, + ) + response.raise_for_status() + return + except httpx.HTTPError: + if time.monotonic() >= deadline: + raise + await asyncio.sleep(1) + + +async def _heartbeat(consul_url: str, worker_id: str) -> None: + check_id = f"service:{worker_id}" + while True: + async with httpx.AsyncClient(timeout=10) as client: + response = await client.put( + f"{consul_url}/v1/agent/check/pass/{check_id}" + ) + response.raise_for_status() + await asyncio.sleep(10) + + +async def _deregister_from_consul(consul_url: str, worker_id: str) -> None: + async with httpx.AsyncClient(timeout=10) as client: + with contextlib.suppress(httpx.HTTPError): + await client.put( + f"{consul_url}/v1/agent/service/deregister/{worker_id}" + ) + + +def create_worker_app( + session_ttl_seconds: int = 900, + session_start_timeout_seconds: int = 30, +) -> FastAPI: + """Build the internal session-worker HTTP application.""" + worker = SessionWorker(session_ttl_seconds, session_start_timeout_seconds) + + consul_url = os.environ.get("CONSUL_HTTP_URL", "http://consul:8500") + consul_url = consul_url.rstrip("/") + worker_id = os.environ.get("WORKER_ID", socket.gethostname()) + worker_address = os.environ.get("WORKER_ADDRESS", socket.gethostname()) + worker_port = int(os.environ.get("WORKER_PORT", "8080")) + worker_endpoint = os.environ.get("WORKER_ENDPOINT") + + @asynccontextmanager + async def lifespan(_app): + await _register_with_consul( + consul_url, + worker_id, + worker_address, + worker_port, + worker_endpoint, + ) + heartbeat = asyncio.create_task(_heartbeat(consul_url, worker_id)) + try: + yield + finally: + heartbeat.cancel() + with contextlib.suppress(asyncio.CancelledError): + await heartbeat + await worker.close_all() + await _deregister_from_consul(consul_url, worker_id) + + app = FastAPI( + title="Select AI Session Worker", + docs_url=None, + redoc_url=None, + lifespan=lifespan, + ) + + @app.get("/health") + async def health() -> dict[str, str]: + return {"status": "ok"} + + @app.post("/sessions", status_code=status.HTTP_201_CREATED) + async def open_session(request: OpenSessionRequest) -> dict[str, str]: + await worker.open(request) + return {"status": "opened"} + + @app.post("/sessions/{session_id}/messages") + async def send_message( + session_id: str, + request: PromptRequest, + ) -> PlainTextResponse: + result = await worker.send_prompt(session_id, request.prompt) + return PlainTextResponse(result or "") + + @app.delete( + "/sessions/{session_id}", + status_code=status.HTTP_204_NO_CONTENT, + ) + async def close_session(session_id: str) -> None: + await worker.close(session_id) + + return app diff --git a/src/select_ai/agent/a2a/worker_client.py b/src/select_ai/agent/a2a/worker_client.py new file mode 100644 index 0000000..b149453 --- /dev/null +++ b/src/select_ai/agent/a2a/worker_client.py @@ -0,0 +1,174 @@ +# ----------------------------------------------------------------------------- +# Copyright (c) 2026, Oracle and/or its affiliates. +# +# Licensed under the Universal Permissive License v 1.0 as shown at +# https://oss.oracle.com/licenses/upl. +# ----------------------------------------------------------------------------- + +"""Consul-backed routing client for the internal session worker.""" + +from __future__ import annotations + +import base64 +import json +import time +from threading import Lock + +import requests + +from select_ai.agent.a2a import GatewaySettings, SessionInfo, SessionRoute + + +class ReconnectRequired(RuntimeError): + """The worker no longer owns the requested in-memory session.""" + + +class WorkerClient: + """Open, route, and close short-lived Select AI worker sessions.""" + + def __init__(self, settings: GatewaySettings): + self.settings = settings + self._selection_lock = Lock() + self._next_worker = 0 + self._worker_request_kwargs: dict[str, object] = {} + if settings.worker_mtls_enabled: + self._worker_request_kwargs = { + "verify": settings.worker_tls_ca_file, + "cert": ( + settings.worker_tls_cert_file, + settings.worker_tls_key_file, + ), + } + + def open_session(self, session_id: str, session_info: SessionInfo) -> str: + """Open a session and save a non-secret route in Consul.""" + endpoint = self._select_worker() + response = requests.post( + f"{endpoint}/sessions", + json={"session_id": session_id, **session_info.__dict__}, + timeout=45, + **self._worker_request_kwargs, + ) + response.raise_for_status() + route = SessionRoute( + endpoint=endpoint, + expires_at=time.time() + self.settings.session_ttl_seconds, + ) + if not self._save_route(session_id, route): + self._close_worker_session(route, session_id) + raise RuntimeError("Could not create the database session.") + return session_id + + def send_prompt(self, session_id: str, prompt: str) -> str | None: + """Forward a prompt and return the worker's raw team result.""" + route = self._route_for(session_id) + response = requests.post( + f"{route.endpoint}/sessions/{session_id}/messages", + json={"prompt": prompt}, + timeout=130, + **getattr(self, "_worker_request_kwargs", {}), + ) + if response.status_code in (404, 502): + self._close_worker_session(route, session_id) + raise ReconnectRequired( + "Database session ended; reconnect required." + ) + response.raise_for_status() + return response.text or None + + def close_session(self, session_id: str) -> None: + """Close the child process and remove the Consul route.""" + try: + route = self._route_for(session_id) + except ReconnectRequired: + self._delete_route(session_id) + return + self._close_worker_session(route, session_id) + + def _select_worker(self) -> str: + response = requests.get( + f"{self.settings.consul_url}/v1/health/service/" + f"{self.settings.worker_service}", + params={"passing": "true"}, + timeout=10, + ) + response.raise_for_status() + workers = response.json() + if not workers: + raise RuntimeError("No healthy Select AI workers are available.") + with self._selection_lock: + worker = workers[self._next_worker % len(workers)] + self._next_worker += 1 + service = worker["Service"] + endpoint = service.get("Meta", {}).get("endpoint") + if endpoint: + endpoint = endpoint.rstrip("/") + if self.settings.worker_mtls_enabled and not endpoint.startswith( + "https://" + ): + raise RuntimeError( + "A worker registered a non-HTTPS endpoint while mTLS is " + "required." + ) + return endpoint + if self.settings.worker_mtls_enabled: + raise RuntimeError( + "Workers must register an HTTPS endpoint while mTLS is " + "required." + ) + address = service.get("Address") or worker["Node"]["Address"] + return f"http://{address}:{service['Port']}" + + def _route_for(self, session_id: str) -> SessionRoute: + response = requests.get( + f"{self.settings.consul_url}/v1/kv/select-ai/sessions/" + f"{session_id}", + timeout=10, + ) + if response.status_code == 404: + raise ReconnectRequired( + "Database session expired; reconnect required." + ) + response.raise_for_status() + value = response.json()[0]["Value"] + route = SessionRoute(**json.loads(base64.b64decode(value).decode())) + if route.expires_at <= time.time(): + self._close_worker_session(route, session_id) + raise ReconnectRequired( + "Database session expired; reconnect required." + ) + return route + + def _save_route(self, session_id: str, route: SessionRoute) -> bool: + response = requests.put( + f"{self.settings.consul_url}/v1/kv/select-ai/sessions/" + f"{session_id}?cas=0", + data=json.dumps(route.__dict__), + timeout=10, + ) + return response.ok and response.text.strip().lower() == "true" + + def _delete_route(self, session_id: str) -> None: + requests.delete( + f"{self.settings.consul_url}/v1/kv/select-ai/sessions/" + f"{session_id}", + timeout=10, + ) + + def _close_worker_session( + self, + route: SessionRoute, + session_id: str, + ) -> None: + try: + response = requests.delete( + f"{route.endpoint}/sessions/{session_id}", + timeout=10, + **getattr(self, "_worker_request_kwargs", {}), + ) + if response.status_code != 404: + response.raise_for_status() + except requests.RequestException: + pass + finally: + self._delete_route(session_id) diff --git a/src/select_ai/cli/a2a.py b/src/select_ai/cli/a2a.py index 9e3bb70..e43342e 100644 --- a/src/select_ai/cli/a2a.py +++ b/src/select_ai/cli/a2a.py @@ -7,6 +7,7 @@ import getpass import json +import ssl import click @@ -92,6 +93,156 @@ def serve( uvicorn.run(app, host=host, port=port) +@a2a.command("worker") +@click.option("--host", default="0.0.0.0", show_default=True) +@click.option("--port", default=8080, show_default=True, type=int) +@click.option( + "--session-ttl-seconds", + default=900, + show_default=True, + type=click.IntRange(min=1), +) +@click.option( + "--session-start-timeout-seconds", + default=30, + show_default=True, + type=click.IntRange(min=1), +) +@click.option( + "--tls-cert-file", + type=click.Path(exists=True, dir_okay=False, readable=True), + help="Worker TLS server certificate. Requires all --tls-* options.", +) +@click.option( + "--tls-key-file", + type=click.Path(exists=True, dir_okay=False, readable=True), + help="Worker TLS server private key. Requires all --tls-* options.", +) +@click.option( + "--tls-ca-file", + type=click.Path(exists=True, dir_okay=False, readable=True), + help="CA used to validate the gateway client certificate.", +) +def worker( + host, + port, + session_ttl_seconds, + session_start_timeout_seconds, + tls_cert_file, + tls_key_file, + tls_ca_file, +): + """Start the internal, in-memory Select AI session worker.""" + try: + from select_ai.agent.a2a import create_worker_app + except ImportError as error: + raise click.ClickException( + "Worker support requires the optional 'a2a' extra. " + "Install it with: pip install 'select_ai[a2a]'" + ) from error + + app = create_worker_app( + session_ttl_seconds=session_ttl_seconds, + session_start_timeout_seconds=session_start_timeout_seconds, + ) + tls_files = (tls_cert_file, tls_key_file, tls_ca_file) + if any(tls_files) and not all(tls_files): + raise click.ClickException( + "Worker mTLS requires --tls-cert-file, --tls-key-file, and " + "--tls-ca-file." + ) + uvicorn_options = {} + if tls_cert_file: + uvicorn_options = { + "ssl_certfile": tls_cert_file, + "ssl_keyfile": tls_key_file, + "ssl_ca_certs": tls_ca_file, + "ssl_cert_reqs": ssl.CERT_REQUIRED, + } + uvicorn.run(app, host=host, port=port, **uvicorn_options) + + +@a2a.command("gateway") +@click.option("--host", default="0.0.0.0", show_default=True) +@click.option("--port", default=8080, show_default=True, type=int) +@click.option( + "--agent-url", + required=True, + envvar="AGENT_URL", + help="Public base URL advertised in the gateway Agent Card.", +) +@click.option( + "--consul-url", + default="http://consul:8500", + show_default=True, + envvar="CONSUL_HTTP_URL", + help="Consul HTTP API URL.", +) +@click.option( + "--worker-service", + default="select-ai-worker", + show_default=True, + envvar="WORKER_SERVICE", + help="Consul service name for Select AI workers.", +) +@click.option( + "--session-ttl-seconds", + default=900, + show_default=True, + type=click.IntRange(min=1), + envvar="SESSION_TTL_SECONDS", +) +@click.option( + "--worker-tls-ca-file", + envvar="WORKER_TLS_CA_FILE", + type=click.Path(exists=True, dir_okay=False, readable=True), + help="CA bundle used to validate worker certificates.", +) +@click.option( + "--worker-tls-cert-file", + envvar="WORKER_TLS_CERT_FILE", + type=click.Path(exists=True, dir_okay=False, readable=True), + help="Gateway client certificate used for worker mTLS.", +) +@click.option( + "--worker-tls-key-file", + envvar="WORKER_TLS_KEY_FILE", + type=click.Path(exists=True, dir_okay=False, readable=True), + help="Gateway client private key used for worker mTLS.", +) +def gateway( + host, + port, + agent_url, + consul_url, + worker_service, + session_ttl_seconds, + worker_tls_ca_file, + worker_tls_cert_file, + worker_tls_key_file, +): + """Start the public A2A/A2UI database-session gateway.""" + try: + from select_ai.agent.a2a import GatewaySettings, create_gateway_app + except ImportError as error: + raise click.ClickException( + "Gateway support requires the optional 'a2a' extra. " + "Install it with: pip install 'select_ai[a2a]'" + ) from error + + settings = GatewaySettings( + agent_url=agent_url, + consul_url=consul_url, + worker_service=worker_service, + session_ttl_seconds=session_ttl_seconds, + worker_tls_ca_file=worker_tls_ca_file, + worker_tls_cert_file=worker_tls_cert_file, + worker_tls_key_file=worker_tls_key_file, + ) + app = create_gateway_app(settings) + uvicorn.run(app, host=host, port=port) + + @a2a.command("agent-card") @click.option("--team", "team_name", required=True, help="Database AI team.") @click.option( diff --git a/tests/a2a/test_agent_card.py b/tests/a2a/test_agent_card.py index 8b543c4..edd594c 100644 --- a/tests/a2a/test_agent_card.py +++ b/tests/a2a/test_agent_card.py @@ -12,6 +12,8 @@ pytest.importorskip("a2a") +from select_ai.agent.a2a.gateway import create_gateway_app +from select_ai.agent.a2a.models import GatewaySettings from select_ai.agent.a2a.server import ( _build_agent_card, _build_v03_agent_card, @@ -31,6 +33,21 @@ def test_v03_discovery_card_is_gemini_enterprise_compatible(): assert payload["protocolVersion"] == "0.3" assert payload["url"] == "https://agent.example.com/a2a/jsonrpc/" assert "supportedInterfaces" not in payload + assert "application/json+a2ui" in payload["defaultOutputModes"] + assert payload["capabilities"]["extensions"] == [ + { + "description": "Provides agent driven UI using the A2UI JSON format.", + "params": { + "acceptsInlineCatalogs": True, + "supportedCatalogIds": [ + "https://www.gstatic.com/vertexaisearch/a2ui/v0_9/" + "gemini_enterprise_composite_catalog.json" + ], + }, + "required": False, + "uri": "https://a2ui.org/a2a-extension/a2ui/v0.9", + } + ] def test_discovery_route_serves_only_the_v03_agent_card(): @@ -53,3 +70,41 @@ def test_discovery_route_serves_only_the_v03_agent_card(): assert payload["protocolVersion"] == "0.3" assert payload["url"] == "https://agent.example.com/a2a/jsonrpc/" assert "supportedInterfaces" not in payload + + +def test_gateway_card_advertises_a2ui_input_and_output(): + app = create_gateway_app( + GatewaySettings( + agent_url="https://agent.example.com", + consul_url="http://consul:8500", + worker_service="select-ai-worker", + session_ttl_seconds=900, + ) + ) + route = next( + route + for route in app.routes + if route.path == "/.well-known/agent-card.json" + ) + + response = asyncio.run(route.endpoint(None)) + payload = json.loads(response.body) + + assert "application/json+a2ui" in payload["defaultInputModes"] + assert "application/json+a2ui" in payload["defaultOutputModes"] + assert "application/json+a2ui" in payload["skills"][0]["inputModes"] + assert "application/json+a2ui" in payload["skills"][0]["outputModes"] + assert payload["capabilities"]["extensions"] == [ + { + "description": "Provides agent driven UI using the A2UI JSON format.", + "params": { + "acceptsInlineCatalogs": True, + "supportedCatalogIds": [ + "https://www.gstatic.com/vertexaisearch/a2ui/v0_9/" + "gemini_enterprise_composite_catalog.json" + ], + }, + "required": False, + "uri": "https://a2ui.org/a2a-extension/a2ui/v0.9", + } + ] diff --git a/tests/a2a/test_worker_runtime.py b/tests/a2a/test_worker_runtime.py new file mode 100644 index 0000000..ce9f953 --- /dev/null +++ b/tests/a2a/test_worker_runtime.py @@ -0,0 +1,382 @@ +# ----------------------------------------------------------------------------- +# Copyright (c) 2026, Oracle and/or its affiliates. +# +# Licensed under the Universal Permissive License v 1.0 as shown at +# https://oss.oracle.com/licenses/upl. +# ----------------------------------------------------------------------------- + +"""Unit tests for the private multiprocessing session runtime.""" + +import asyncio + +import pytest + +pytest.importorskip("fastapi") + +import select_ai +from select_ai.agent.a2a import GatewaySettings, worker +from select_ai.agent.a2a.gateway import GatewayExecutor, _a2ui_part +from select_ai.agent.a2a.results import message_parts +from select_ai.agent.a2a.worker_client import WorkerClient + + +class FakeConnection: + """Small synchronous pipe stand-in for worker protocol tests.""" + + def __init__(self, responses=()): + self.responses = list(responses) + self.sent = [] + self.closed = False + + def send(self, value): + self.sent.append(value) + + def poll(self, _timeout): + return bool(self.responses) + + def recv(self): + return self.responses.pop(0) + + def close(self): + self.closed = True + + +class FakeProcess: + """Process stand-in which lets tests assert IPC without forking.""" + + def __init__(self, *, target, args, daemon): + self.target = target + self.args = args + self.daemon = daemon + self.started = False + self.alive = True + + def start(self): + self.started = True + + def is_alive(self): + return self.alive + + def join(self, timeout): + del timeout + + def terminate(self): + self.alive = False + + +def test_worker_uses_pipe_process_and_returns_raw_result(monkeypatch): + parent = FakeConnection( + [{"type": "ready"}, {"type": "result", "result": "hi"}] + ) + child = FakeConnection() + processes = [] + + def process_factory(**kwargs): + process = FakeProcess(**kwargs) + processes.append(process) + return process + + monkeypatch.setattr( + worker.multiprocessing, + "Pipe", + lambda: (parent, child), + ) + monkeypatch.setattr(worker.multiprocessing, "Process", process_factory) + session_worker = worker.SessionWorker(60, 1) + request = worker.OpenSessionRequest( + session_id="session-1", + username="user", + password="password", + dsn="database", + team_name="TEAM", + ) + + asyncio.run(session_worker.open(request)) + result = asyncio.run(session_worker.send_prompt("session-1", "hello")) + + assert result == "hi" + assert child.closed + assert processes[0].target is worker._session_process_main + assert processes[0].daemon is True + assert parent.sent == [{"type": "run", "prompt": "hello"}] + + +def test_expired_session_closes_the_process_and_pipe(monkeypatch): + parent = FakeConnection([{"type": "ready"}]) + child = FakeConnection() + process = None + + def process_factory(**kwargs): + nonlocal process + process = FakeProcess(**kwargs) + return process + + monkeypatch.setattr( + worker.multiprocessing, + "Pipe", + lambda: (parent, child), + ) + monkeypatch.setattr(worker.multiprocessing, "Process", process_factory) + session_worker = worker.SessionWorker(60, 1) + request = worker.OpenSessionRequest( + session_id="session-1", + username="user", + password="password", + dsn="database", + team_name="TEAM", + ) + asyncio.run(session_worker.open(request)) + session_worker.sessions["session-1"].expires_at = 0 + + with pytest.raises(worker.HTTPException, match="reconnect required"): + asyncio.run(session_worker.get("session-1")) + + assert process is not None + assert process.alive is False + assert parent.closed + assert parent.sent == [{"type": "close"}] + + +def test_session_runtime_uses_one_async_connection_and_returns_raw_result( + monkeypatch, +): + connection = FakeConnection( + [{"type": "run", "prompt": "hello"}, {"type": "close"}] + ) + connection_arguments = {} + + class Conversation: + def __init__(self, attributes): + self.attributes = attributes + + async def create(self): + return "conversation-1" + + class Team: + def __init__(self, team_name): + assert team_name == "TEAM" + + async def run(self, prompt, params): + assert prompt == "hello" + assert params == {"conversation_id": "conversation-1"} + return ( + '{"metadata":{"mimeType":"application/json+a2ui"},' + '"data":[]}' + ) + + async def async_connect(**kwargs): + connection_arguments.update(kwargs) + + async def connected(): + return True + + async def disconnect(): + return None + + monkeypatch.setattr(select_ai, "async_connect", async_connect) + monkeypatch.setattr(select_ai, "async_is_connected", connected) + monkeypatch.setattr(select_ai, "async_disconnect", disconnect) + monkeypatch.setattr(select_ai, "AsyncConversation", Conversation) + monkeypatch.setattr(select_ai.agent, "AsyncTeam", Team) + + asyncio.run( + worker._run_session_process( + connection, + { + "user": "user", + "password": "password", + "dsn": "database", + }, + "session-1", + "TEAM", + ) + ) + + assert connection_arguments == { + "user": "user", + "password": "password", + "dsn": "database", + } + assert connection.sent[0] == {"type": "ready"} + assert connection.sent[1] == { + "type": "result", + "result": ( + '{"metadata":{"mimeType":"application/json+a2ui"},' '"data":[]}' + ), + } + + +def test_message_parts_preserves_serialized_text_and_data_parts(): + parts = message_parts( + """{ + "kind": "message", + "parts": [ + {"kind": "text", "text": "hello"}, + { + "kind": "data", + "data": {"value": 1}, + "metadata": {"mimeType": "application/json+a2ui"} + } + ] + }""" + ) + + assert parts is not None + assert parts[0].text == "hello" + assert parts[1].WhichOneof("content") == "data" + + +def test_worker_client_returns_the_raw_team_result(monkeypatch): + class Response: + status_code = 200 + + @staticmethod + def raise_for_status(): + return None + + text = "team result" + + monkeypatch.setattr( + "select_ai.agent.a2a.worker_client.requests.post", + lambda *args, **kwargs: Response(), + ) + client = WorkerClient.__new__(WorkerClient) + monkeypatch.setattr( + client, + "_route_for", + lambda session_id: type("Route", (), {"endpoint": "http://worker"})(), + ) + + assert client.send_prompt("context-1", "hello") == "team result" + + +def test_worker_client_uses_consul_https_endpoint_with_mtls(monkeypatch): + settings = GatewaySettings( + agent_url="https://gateway.example.com", + consul_url="http://consul:8500", + worker_service="select-ai-worker", + session_ttl_seconds=60, + worker_tls_ca_file="/tls/ca.pem", + worker_tls_cert_file="/tls/gateway.pem", + worker_tls_key_file="/tls/gateway-key.pem", + ) + client = WorkerClient(settings) + + class Response: + @staticmethod + def raise_for_status(): + return None + + @staticmethod + def json(): + return [ + { + "Service": { + "Meta": {"endpoint": "https://worker-0.internal"} + }, + "Node": {}, + } + ] + + monkeypatch.setattr( + "select_ai.agent.a2a.worker_client.requests.get", + lambda *args, **kwargs: Response(), + ) + + assert client._select_worker() == "https://worker-0.internal" + assert client._worker_request_kwargs == { + "verify": "/tls/ca.pem", + "cert": ("/tls/gateway.pem", "/tls/gateway-key.pem"), + } + + +def test_mtls_requires_all_three_gateway_files(): + with pytest.raises(ValueError, match="worker mTLS requires"): + GatewaySettings( + agent_url="https://gateway.example.com", + consul_url="http://consul:8500", + worker_service="select-ai-worker", + session_ttl_seconds=60, + worker_tls_ca_file="/tls/ca.pem", + ) + + +def test_worker_registers_its_https_endpoint(monkeypatch): + registered = {} + + class Response: + @staticmethod + def raise_for_status(): + return None + + class Client: + def __init__(self, **kwargs): + del kwargs + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + return None + + async def put(self, _url, json): + registered.update(json) + return Response() + + monkeypatch.setattr(worker.httpx, "AsyncClient", Client) + asyncio.run( + worker._register_with_consul( + "http://consul:8500", + "worker-0", + "10.0.0.1", + 8080, + "https://worker-0.internal", + ) + ) + + assert registered["Meta"] == {"endpoint": "https://worker-0.internal"} + + +def test_a2ui_operation_uses_a_metadata_marked_data_part(): + from google.protobuf.json_format import MessageToDict + + part = _a2ui_part( + {"version": "v0.9", "createSurface": {"surfaceId": "form"}} + ) + + assert MessageToDict(part.data) == { + "version": "v0.9", + "createSurface": {"surfaceId": "form"}, + } + assert MessageToDict(part.metadata) == { + "mimeType": "application/json+a2ui" + } + + +def test_a2ui_action_reads_an_operation_list_or_single_operation(): + action = GatewayExecutor._a2ui_action( + type( + "Context", + (), + { + "message": type( + "Message", + (), + { + "parts": [ + _a2ui_part( + { + "version": "v0.9", + "action": { + "name": "submit_database_connection" + }, + } + ) + ] + }, + )() + }, + )() + ) + + assert action == {"name": "submit_database_connection"} From 1f42af0e52a9680e03e94c6ee337e67528753937 Mon Sep 17 00:00:00 2001 From: Abhishek Singh Date: Mon, 31 Aug 2026 20:44:38 -0700 Subject: [PATCH 2/2] Fixes after E2E testing with Gemini Enterprise --- gcloud/gateway/README.md | 116 ++++++++++++++----------- gcloud/gateway/cloudbuild.yaml | 48 +++++----- gcloud/gateway/deploy.sh | 103 ++++++++++++++++------ gcloud/gateway/gke/worker-mtls.yaml | 4 +- gcloud/gateway/gke/worker-service.yaml | 18 ++++ gcloud/gateway/gke/worker.yaml | 2 + src/select_ai/agent/a2a/gateway.py | 3 +- tests/a2a/test_worker_runtime.py | 28 ++++++ 8 files changed, 216 insertions(+), 106 deletions(-) create mode 100644 gcloud/gateway/gke/worker-service.yaml diff --git a/gcloud/gateway/README.md b/gcloud/gateway/README.md index a992159..a01ee79 100644 --- a/gcloud/gateway/README.md +++ b/gcloud/gateway/README.md @@ -1,21 +1,33 @@ # Dynamic gateway deployment -Dynamic gateway mode exposes one public A2A endpoint where each user selects -an Oracle database connection and Select AI team through the A2UI connection -form. +Dynamic gateway mode exposes one public A2A endpoint. Each user dynamically +selects an Oracle database connection and Select AI team through the A2UI +connection form. A session remains available for 15 minutes by default. Set a +different lifetime in seconds with `--session-ttl-seconds`; for example, +`--session-ttl-seconds 1800` keeps sessions for 30 minutes. ```text -A2A client - | -Cloud Run gateway - | - +-- private VPC --> Consul in GKE - | - +-- private VPC --> worker replicas in GKE - | - session child process - | - Oracle Database + ┌──────────────────────┐ + │ A2A / Gemini client │ + └──────────┬───────────┘ + │ public A2A + v + ┌──────────────────────┐ + │ Cloud Run gateway │ + └──────┬───────┬───────┘ + │ │ private VPC: mTLS request to worker hostname + │ │ + │ │ ┌─────────────────────── GKE ───────────────────────┐ + │ └────>│ [Headless Service + managed VPC DNS] │ + │ │ worker hostname → current worker Pod IP │ + │ │ │ │ + │ │ v │ + │ │ [StatefulSet worker-0 / worker-1 / ...] │ + │ │ session child process → Oracle Database │ + │ │ │ + │ │ [Consul] │ + └────────────>│ selects healthy worker; returns worker hostname │ + └───────────────────────────────────────────────────┘ ``` The gateway is the only public A2A application. Consul and workers are a GKE @@ -31,7 +43,9 @@ gcloud/gateway/deploy.sh --project PROJECT_ID ``` The script creates the Artifact Registry repository and GKE Autopilot cluster -when they do not already exist. Cloud Build then: +when they do not already exist. The cluster is created with +GKE additive VPC DNS: GKE owns the worker DNS records and keeps them current +when a worker Pod is recreated. Cloud Build then: 1. builds the existing `docker/Dockerfile` image once; 2. deploys the GKE namespace and internal Consul service; @@ -47,6 +61,7 @@ gcloud/gateway/deploy.sh \ --project PROJECT_ID \ --region us-central1 \ --cluster select-ai-a2a-gateway \ + --gke-dns-domain select-ai-a2a-gateway.internal \ --worker-replicas 3 \ --network default \ --subnet default @@ -67,13 +82,12 @@ deployment time. Local testing does not use mTLS. The default GCloud deployment also keeps the current private-VPC HTTP worker transport. -For a short-lived GCloud mTLS test, add a private DNS suffix: +For a short-lived GCloud mTLS test: ```bash gcloud/gateway/deploy.sh \ --project PROJECT_ID \ --enable-worker-mtls \ - --worker-domain workers.select-ai.internal \ --mtls-cert-validity-days 365 ``` @@ -93,31 +107,45 @@ reuse them, including when changing `--worker-replicas`. To deliberately replace the CA and both leaf certificates, add `--rotate-worker-mtls`. Rotation recreates the worker StatefulSet and ends active worker sessions. -Workers run as a StatefulSet. Each worker registers its own private DNS name -with Consul, so the gateway still reaches the exact process that owns the -session. Cloud Build creates the corresponding private Cloud DNS A records. +Workers run as a StatefulSet. The `select-ai-worker` headless Service gives +each worker a stable name, for example +`select-ai-worker-0.select-ai-worker.select-ai-gateway.svc.select-ai-a2a-gateway.internal`. +Consul registers that name, so the gateway reaches the exact worker that owns a +session. GKE Cloud DNS updates its Pod-IP record automatically after a worker +is recreated. There is no worker load balancer, custom Cloud DNS zone, or +deployment-time Pod-IP snapshot. -The certificate and DNS records are refreshed by rerunning the deployment. -Because this is a test mode, rerun it after a worker pod is recreated outside a -deployment; its pod IP can change. Do not use this mode for a long-lived -production stack—use a managed workload-certificate and DNS reconciliation -solution there. +`--gke-dns-domain` must be unique in the VPC and cannot end in `.local`. It is +immutable after cluster creation. Autopilot supports additive VPC DNS only when +the cluster is created, so an older cluster without it cannot be reused by this +deployment. Use a new `--cluster` name for the first migration, verify it, then +delete the old cluster when you are ready. -### Where the worker certificate files come from +If existing mTLS material was issued for a different GKE DNS domain, the script +replaces it automatically before deploying the replacement cluster. -The certificate paths passed to `select-ai a2a worker` are files **inside each -GKE worker container**. They are mounted from Kubernetes Secrets, not files in -this repository or on the machine that runs `deploy.sh`. +### How a worker DNS name is decided + +GKE gives a StatefulSet Pod a DNS name using this form: + +```text +...svc. +``` + +For this deployment, worker 0 is: ```text -deploy.sh - -> creates the test certificates and stores them in Secret Manager - -> Cloud Build reads the worker certificate, worker key, and CA certificate - -> Cloud Build creates Kubernetes Secrets in select-ai-gateway - -> GKE mounts those Secrets read-only in every worker container +select-ai-worker-0.select-ai-worker.select-ai-gateway.svc.select-ai-a2a-gateway.internal ``` -The worker pod receives these mounts: +`select-ai-worker-0` is the StatefulSet Pod name, `select-ai-worker` is the +headless Service, `select-ai-gateway` is the Kubernetes namespace, and +`select-ai-a2a-gateway.internal` is the `--gke-dns-domain` value. GKE updates +the resulting record when the Pod IP changes. + +### Certificate mounts + +Worker certificate files are mounted from Kubernetes Secrets: | Container file | Kubernetes Secret | Secret key | Used for | | --- | --- | --- | --- | @@ -125,21 +153,7 @@ The worker pod receives these mounts: | `/var/run/select-ai-mtls/tls.key` | `select-ai-worker-server-tls` | `tls.key` | worker HTTPS private key | | `/var/run/select-ai-mtls/gateway-ca.crt` | `select-ai-gateway-client-ca` | `ca.crt` | validates the gateway client certificate | -GKE does not read Google Secret Manager at worker request time. Cloud Build -copies the required material into Kubernetes Secrets during deployment, and the -worker certificate files are mounted from those Kubernetes Secrets as read-only -files. For example, to inspect the paths in a running worker (without printing -their contents): - -```bash -kubectl -n select-ai-gateway exec select-ai-worker-0 -- \ - ls -l /var/run/select-ai-mtls/ -``` - -The Cloud Run gateway uses a separate direct mount from Google Secret Manager -for its CA, client certificate, and client key. +The Cloud Run gateway certificate files are mounted from Google Secret Manager. The identity that submits Cloud Build needs permission to use GKE, Cloud Run, -Cloud DNS, and Secret Manager. Cloud Build also needs the corresponding GKE, -Cloud Run, Cloud DNS, and Secret Manager permissions because it deploys the -workers, refreshes the private DNS records, and mounts the gateway secrets. +and Secret Manager. GKE maintains the managed worker DNS records. diff --git a/gcloud/gateway/cloudbuild.yaml b/gcloud/gateway/cloudbuild.yaml index 69958b7..89bb0af 100644 --- a/gcloud/gateway/cloudbuild.yaml +++ b/gcloud/gateway/cloudbuild.yaml @@ -25,9 +25,23 @@ steps: - --location=${_REGION} - --cluster=${_CLUSTER} - # gke-deploy writes generated manifests to output/. Cloud Build shares the - # source workspace across steps, so remove only that builder output before - # invoking it again. + # gke-deploy writes generated manifests to output/. Remove its output before + # running it for the headless worker Service. + - name: gcr.io/cloud-builders/gcloud + entrypoint: bash + args: + - -ceu + - rm -rf output + + - name: gcr.io/cloud-builders/gke-deploy + args: + - run + - --filename=gcloud/gateway/gke/worker-service.yaml + - --location=${_REGION} + - --cluster=${_CLUSTER} + + # Cloud Build shares the source workspace across steps, so remove only the + # gke-deploy builder output before invoking another deployment tool. - name: gcr.io/cloud-builders/gcloud entrypoint: bash args: @@ -115,7 +129,8 @@ steps: fi sed -e "s|WORKER_IMAGE|$image|" \ -e "s|WORKER_REPLICAS|${_WORKER_REPLICAS}|" \ - -e "s|WORKER_DOMAIN|${_WORKER_DOMAIN}|" \ + -e "s|SESSION_TTL_SECONDS|${_SESSION_TTL_SECONDS}|" \ + -e "s|GKE_DNS_DOMAIN|${_GKE_DNS_DOMAIN}|" \ "$worker_manifest" > /workspace/worker.yaml - name: gcr.io/cloud-builders/gcloud @@ -141,29 +156,10 @@ steps: - | if [[ "${_ENABLE_WORKER_MTLS}" == "true" ]]; then kubectl rollout status --namespace select-ai-gateway statefulset/select-ai-worker --timeout=300s - kubectl get pods --namespace select-ai-gateway -l app=select-ai-worker \ - -o jsonpath='{range .items[*]}{.metadata.name}={.status.podIP}{"\n"}{end}' \ - > /workspace/worker-pods else kubectl rollout status --namespace select-ai-gateway deployment/select-ai-worker --timeout=300s fi - - name: gcr.io/cloud-builders/gcloud - entrypoint: bash - args: - - -ceu - - | - if [[ "${_ENABLE_WORKER_MTLS}" != "true" ]]; then - exit 0 - fi - while IFS== read -r pod_name pod_ip; do - record_name="${pod_name}.${_WORKER_DOMAIN}." - gcloud dns record-sets delete "$record_name" --type=A \ - --zone="${_WORKER_DNS_ZONE}" --quiet || true - gcloud dns record-sets create "$record_name" --type=A --ttl=30 \ - --rrdatas="$pod_ip" --zone="${_WORKER_DNS_ZONE}" - done < /workspace/worker-pods - - name: gcr.io/cloud-builders/gcloud entrypoint: bash args: @@ -191,7 +187,7 @@ steps: --port=8080 \ --command=select-ai \ --args=a2a,gateway \ - --set-env-vars="AGENT_URL=https://pending.invalid,CONSUL_HTTP_URL=http://$consul_ip:8500,WORKER_SERVICE=select-ai-worker,SESSION_TTL_SECONDS=900$mtls_env" \ + --set-env-vars="AGENT_URL=https://pending.invalid,CONSUL_HTTP_URL=http://$consul_ip:8500,WORKER_SERVICE=select-ai-worker,SESSION_TTL_SECONDS=${_SESSION_TTL_SECONDS}$mtls_env" \ "${mtls_args[@]}" gateway_url=$(gcloud run services describe "${_GATEWAY_SERVICE}" \ --region="${_REGION}" --format='value(status.url)') @@ -211,10 +207,10 @@ substitutions: _NETWORK: default _SUBNET: default _WORKER_REPLICAS: "2" + _SESSION_TTL_SECONDS: "900" _ENABLE_WORKER_MTLS: "false" _ROTATE_WORKER_MTLS: "false" - _WORKER_DOMAIN: "" - _WORKER_DNS_ZONE: select-ai-workers + _GKE_DNS_DOMAIN: select-ai-a2a-gateway.internal options: logging: CLOUD_LOGGING_ONLY diff --git a/gcloud/gateway/deploy.sh b/gcloud/gateway/deploy.sh index 1923940..55959c8 100755 --- a/gcloud/gateway/deploy.sh +++ b/gcloud/gateway/deploy.sh @@ -18,17 +18,19 @@ Usage: gcloud/gateway/deploy.sh [options] Options: --project PROJECT Google Cloud project (defaults to gcloud config) - --region REGION Region for GKE, Cloud Run, and Artifact Registry (default: us-central1) + --region REGION Region for GKE, Cloud Run, and Artifact Registry + (default: us-central1) --cluster NAME GKE Autopilot cluster name (default: select-ai-a2a-gateway) + --gke-dns-domain DOMAIN Unique GKE additive VPC DNS domain + (default: select-ai-a2a-gateway.internal) --repository NAME Artifact Registry Docker repository (default: select-ai) --gateway-service NAME Cloud Run gateway service name (default: select-ai-a2a-gateway) --network NAME VPC network for Cloud Run direct VPC egress (default: default) --subnet NAME VPC subnet for Cloud Run direct VPC egress (default: default) --worker-replicas COUNT GKE worker replica count (default: 2) + --session-ttl-seconds N Dynamic session lifetime (default: 900) --enable-worker-mtls Use ephemeral mTLS certificates for gateway-to-worker calls --rotate-worker-mtls Replace the existing worker mTLS CA and certificates - --worker-domain DOMAIN Private DNS suffix for workers; required with --enable-worker-mtls - --worker-dns-zone NAME Private Cloud DNS zone name (default: select-ai-workers) --mtls-cert-validity-days DAYS Gateway and worker certificate lifetime (default: 365) --image-tag TAG Image tag (default: git SHA plus UTC timestamp) @@ -40,15 +42,15 @@ repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" project_id="" region="us-central1" cluster="select-ai-a2a-gateway" +gke_dns_domain="select-ai-a2a-gateway.internal" repository="select-ai" gateway_service="select-ai-a2a-gateway" network="default" subnet="default" worker_replicas="2" +session_ttl_seconds="900" enable_worker_mtls="false" rotate_worker_mtls="false" -worker_domain="" -worker_dns_zone="select-ai-workers" mtls_cert_validity_days="365" image_tag="" @@ -57,16 +59,22 @@ while [[ $# -gt 0 ]]; do --project) project_id="${2:?--project requires a value}"; shift 2 ;; --region) region="${2:?--region requires a value}"; shift 2 ;; --cluster) cluster="${2:?--cluster requires a value}"; shift 2 ;; + --gke-dns-domain) gke_dns_domain="${2:?--gke-dns-domain requires a value}"; shift 2 ;; --repository) repository="${2:?--repository requires a value}"; shift 2 ;; --gateway-service) gateway_service="${2:?--gateway-service requires a value}"; shift 2 ;; --network) network="${2:?--network requires a value}"; shift 2 ;; --subnet) subnet="${2:?--subnet requires a value}"; shift 2 ;; --worker-replicas) worker_replicas="${2:?--worker-replicas requires a value}"; shift 2 ;; + --session-ttl-seconds) + session_ttl_seconds="${2:?--session-ttl-seconds requires a value}" + shift 2 + ;; --enable-worker-mtls) enable_worker_mtls="true"; shift ;; --rotate-worker-mtls) rotate_worker_mtls="true"; shift ;; - --worker-domain) worker_domain="${2:?--worker-domain requires a value}"; shift 2 ;; - --worker-dns-zone) worker_dns_zone="${2:?--worker-dns-zone requires a value}"; shift 2 ;; - --mtls-cert-validity-days) mtls_cert_validity_days="${2:?--mtls-cert-validity-days requires a value}"; shift 2 ;; + --mtls-cert-validity-days) + mtls_cert_validity_days="${2:?--mtls-cert-validity-days requires a value}" + shift 2 + ;; --image-tag) image_tag="${2:?--image-tag requires a value}"; shift 2 ;; -h|--help) usage; exit 0 ;; *) echo "Unknown option: $1" >&2; usage >&2; exit 2 ;; @@ -84,8 +92,12 @@ if ! [[ "$worker_replicas" =~ ^[1-9][0-9]*$ ]]; then echo "--worker-replicas must be a positive integer." >&2 exit 2 fi -if [[ "$enable_worker_mtls" == "true" && -z "$worker_domain" ]]; then - echo "--worker-domain is required with --enable-worker-mtls." >&2 +if ! [[ "$session_ttl_seconds" =~ ^[1-9][0-9]*$ ]]; then + echo "--session-ttl-seconds must be a positive integer." >&2 + exit 2 +fi +if [[ -z "$gke_dns_domain" || "$gke_dns_domain" == *.local ]]; then + echo "--gke-dns-domain must be non-empty and must not end in .local." >&2 exit 2 fi if [[ "$rotate_worker_mtls" == "true" && "$enable_worker_mtls" != "true" ]]; then @@ -119,15 +131,29 @@ fi if ! gcloud container clusters describe "$cluster" \ --location="$region" --project="$project_id" >/dev/null 2>&1; then gcloud container clusters create-auto "$cluster" \ - --location="$region" --network="$network" --project="$project_id" + --location="$region" --network="$network" --subnetwork="$subnet" \ + --additive-vpc-scope-dns-domain="$gke_dns_domain" --project="$project_id" +else + configured_dns_domain="$(gcloud container clusters describe "$cluster" \ + --location="$region" --project="$project_id" \ + --format='value(networkConfig.dnsConfig.additiveVpcScopeDnsDomain)')" + if [[ "$configured_dns_domain" != "$gke_dns_domain" ]]; then + echo "Cluster $cluster does not use additive VPC DNS domain $gke_dns_domain." >&2 + echo "Create a replacement cluster or pass its configured --gke-dns-domain." >&2 + exit 2 + fi fi +project_number="$(gcloud projects describe "$project_id" \ + --format='value(projectNumber)')" + cleanup_mtls() { [[ -n "${mtls_dir:-}" ]] && rm -rf "$mtls_dir" } if [[ "$enable_worker_mtls" == "true" ]]; then create_mtls_material="$rotate_worker_mtls" + worker_certificate_dns_name="*.select-ai-worker.select-ai-gateway.svc.$gke_dns_domain" for secret_name in \ select-ai-gateway-mtls-ca \ select-ai-gateway-mtls-cert \ @@ -139,6 +165,16 @@ if [[ "$enable_worker_mtls" == "true" ]]; then create_mtls_material="true" fi done + if [[ "$create_mtls_material" != "true" ]]; then + existing_worker_certificate_sans="$(gcloud secrets versions access latest \ + --secret=select-ai-worker-mtls-cert --project="$project_id" 2>/dev/null | \ + openssl x509 -noout -ext subjectAltName 2>/dev/null || true)" + if [[ "$existing_worker_certificate_sans" != *"DNS:$worker_certificate_dns_name"* ]]; then + create_mtls_material="true" + echo "Replacing mTLS material because the worker certificate does not" + echo "match the GKE DNS domain." + fi + fi if [[ "$create_mtls_material" == "true" ]]; then trap cleanup_mtls EXIT mtls_dir="$(mktemp -d)" @@ -149,8 +185,8 @@ if [[ "$enable_worker_mtls" == "true" ]]; then -subj "/CN=select-ai ephemeral worker CA" >/dev/null 2>&1 openssl req -newkey rsa:2048 -nodes \ -keyout "$mtls_dir/worker.key" -out "$mtls_dir/worker.csr" \ - -subj "/CN=*.$worker_domain" >/dev/null 2>&1 - printf 'subjectAltName=DNS:*.%s\nextendedKeyUsage=serverAuth\n' "$worker_domain" \ + -subj "/CN=select-ai-worker" >/dev/null 2>&1 + printf 'subjectAltName=DNS:%s\nextendedKeyUsage=serverAuth\n' "$worker_certificate_dns_name" \ > "$mtls_dir/worker.ext" openssl x509 -req -days "$mtls_cert_validity_days" -in "$mtls_dir/worker.csr" \ -CA "$mtls_dir/ca.crt" -CAkey "$mtls_dir/ca.key" -CAcreateserial \ @@ -171,11 +207,16 @@ if [[ "$enable_worker_mtls" == "true" ]]; then gcloud secrets describe "$secret_name" --project="$project_id" >/dev/null 2>&1 || \ gcloud secrets create "$secret_name" --replication-policy=automatic --project="$project_id" done - gcloud secrets versions add select-ai-gateway-mtls-ca --data-file="$mtls_dir/ca.crt" --project="$project_id" - gcloud secrets versions add select-ai-gateway-mtls-cert --data-file="$mtls_dir/gateway.crt" --project="$project_id" - gcloud secrets versions add select-ai-gateway-mtls-key --data-file="$mtls_dir/gateway.key" --project="$project_id" - gcloud secrets versions add select-ai-worker-mtls-cert --data-file="$mtls_dir/worker.crt" --project="$project_id" - gcloud secrets versions add select-ai-worker-mtls-key --data-file="$mtls_dir/worker.key" --project="$project_id" + gcloud secrets versions add select-ai-gateway-mtls-ca \ + --data-file="$mtls_dir/ca.crt" --project="$project_id" + gcloud secrets versions add select-ai-gateway-mtls-cert \ + --data-file="$mtls_dir/gateway.crt" --project="$project_id" + gcloud secrets versions add select-ai-gateway-mtls-key \ + --data-file="$mtls_dir/gateway.key" --project="$project_id" + gcloud secrets versions add select-ai-worker-mtls-cert \ + --data-file="$mtls_dir/worker.crt" --project="$project_id" + gcloud secrets versions add select-ai-worker-mtls-key \ + --data-file="$mtls_dir/worker.key" --project="$project_id" if [[ "$rotate_worker_mtls" == "true" ]]; then echo "Rotating worker mTLS material; worker sessions will be interrupted." else @@ -184,13 +225,7 @@ if [[ "$enable_worker_mtls" == "true" ]]; then else echo "Reusing existing worker mTLS material." fi - if ! gcloud dns managed-zones describe "$worker_dns_zone" --project="$project_id" >/dev/null 2>&1; then - gcloud dns managed-zones create "$worker_dns_zone" --dns-name="${worker_domain}." \ - --visibility=private --networks="https://www.googleapis.com/compute/v1/projects/$project_id/global/networks/$network" \ - --description="Private DNS records for Select AI mTLS workers" \ - --project="$project_id" - fi - runtime_sa="${project_id_number:-$(gcloud projects describe "$project_id" --format='value(projectNumber)')}-compute@developer.gserviceaccount.com" + runtime_sa="${project_number}-compute@developer.gserviceaccount.com" for secret_name in \ select-ai-gateway-mtls-ca \ select-ai-gateway-mtls-cert \ @@ -202,15 +237,29 @@ if [[ "$enable_worker_mtls" == "true" ]]; then done fi +build_substitutions=( + "_REGION=$region" + "_CLUSTER=$cluster" + "_REPOSITORY=$repository" + "_IMAGE_TAG=$image_tag" + "_GATEWAY_SERVICE=$gateway_service" + "_NETWORK=$network" + "_SUBNET=$subnet" + "_WORKER_REPLICAS=$worker_replicas" + "_SESSION_TTL_SECONDS=$session_ttl_seconds" + "_ENABLE_WORKER_MTLS=$enable_worker_mtls" + "_ROTATE_WORKER_MTLS=$rotate_worker_mtls" + "_GKE_DNS_DOMAIN=$gke_dns_domain" +) + gcloud builds submit "$repo_root" \ --project="$project_id" \ --region="$region" \ --config="$repo_root/gcloud/gateway/cloudbuild.yaml" \ - --substitutions="_REGION=$region,_CLUSTER=$cluster,_REPOSITORY=$repository,_IMAGE_TAG=$image_tag,_GATEWAY_SERVICE=$gateway_service,_NETWORK=$network,_SUBNET=$subnet,_WORKER_REPLICAS=$worker_replicas,_ENABLE_WORKER_MTLS=$enable_worker_mtls,_ROTATE_WORKER_MTLS=$rotate_worker_mtls,_WORKER_DOMAIN=$worker_domain,_WORKER_DNS_ZONE=$worker_dns_zone" + --substitutions="$(IFS=,; printf '%s' "${build_substitutions[*]}")" gateway_url="$(gcloud run services describe "$gateway_service" \ --region="$region" --project="$project_id" --format='value(status.url)')" -project_number="$(gcloud projects describe "$project_id" --format='value(projectNumber)')" gemini_service_agent="service-$project_number@gcp-sa-discoveryengine.iam.gserviceaccount.com" gcloud run services add-iam-policy-binding "$gateway_service" \ diff --git a/gcloud/gateway/gke/worker-mtls.yaml b/gcloud/gateway/gke/worker-mtls.yaml index 9861212..bdab56e 100644 --- a/gcloud/gateway/gke/worker-mtls.yaml +++ b/gcloud/gateway/gke/worker-mtls.yaml @@ -26,6 +26,8 @@ spec: - 0.0.0.0 - --port - "8443" + - --session-ttl-seconds + - "SESSION_TTL_SECONDS" - --tls-cert-file - /var/run/select-ai-mtls/tls.crt - --tls-key-file @@ -53,7 +55,7 @@ spec: - name: WORKER_PORT value: "8443" - name: WORKER_ENDPOINT - value: https://$(POD_NAME).WORKER_DOMAIN:8443 + value: https://$(POD_NAME).select-ai-worker.select-ai-gateway.svc.GKE_DNS_DOMAIN:8443 volumeMounts: - name: worker-server-tls mountPath: /var/run/select-ai-mtls/tls.crt diff --git a/gcloud/gateway/gke/worker-service.yaml b/gcloud/gateway/gke/worker-service.yaml new file mode 100644 index 0000000..103fe2a --- /dev/null +++ b/gcloud/gateway/gke/worker-service.yaml @@ -0,0 +1,18 @@ +apiVersion: v1 +kind: Service +metadata: + name: select-ai-worker + namespace: select-ai-gateway +spec: + # This is service discovery only: it allocates no load balancer and no + # virtual Service IP. Cloud DNS for GKE publishes the current Pod IPs. + clusterIP: None + selector: + app: select-ai-worker + ports: + - name: http + port: 8080 + targetPort: 8080 + - name: https + port: 8443 + targetPort: 8443 diff --git a/gcloud/gateway/gke/worker.yaml b/gcloud/gateway/gke/worker.yaml index fc24c5e..3b5392b 100644 --- a/gcloud/gateway/gke/worker.yaml +++ b/gcloud/gateway/gke/worker.yaml @@ -26,6 +26,8 @@ spec: - 0.0.0.0 - --port - "8080" + - --session-ttl-seconds + - "SESSION_TTL_SECONDS" ports: - name: http containerPort: 8080 diff --git a/src/select_ai/agent/a2a/gateway.py b/src/select_ai/agent/a2a/gateway.py index 4b0957e..622e2c7 100644 --- a/src/select_ai/agent/a2a/gateway.py +++ b/src/select_ai/agent/a2a/gateway.py @@ -11,6 +11,7 @@ import asyncio +import requests from a2a.compat.v0_3.conversions import to_compat_agent_card from a2a.helpers import ( new_data_part, @@ -100,7 +101,7 @@ async def _open_session(self, action_context: dict, context_id: str): context_id, session_info, ) - except (ValueError, RuntimeError): + except (requests.RequestException, ValueError, RuntimeError): return [ new_text_part( "Could not connect. Check the DSN, credentials, and team " diff --git a/tests/a2a/test_worker_runtime.py b/tests/a2a/test_worker_runtime.py index ce9f953..8ab1cff 100644 --- a/tests/a2a/test_worker_runtime.py +++ b/tests/a2a/test_worker_runtime.py @@ -10,6 +10,7 @@ import asyncio import pytest +import requests pytest.importorskip("fastapi") @@ -380,3 +381,30 @@ def test_a2ui_action_reads_an_operation_list_or_single_operation(): ) assert action == {"name": "submit_database_connection"} + + +def test_gateway_returns_connection_error_when_worker_rejects_opening(): + executor = GatewayExecutor.__new__(GatewayExecutor) + executor.sessions = {} + + class Client: + @staticmethod + def open_session(_context_id, _session_info): + raise requests.HTTPError("worker rejected the connection") + + executor.worker_client = Client() + + parts = asyncio.run( + executor._open_session( + { + "dsn": "database", + "username": "user", + "password": "password", + "team_name": "TEAM", + }, + "context-1", + ) + ) + + assert parts[0].text.startswith("Could not connect.") + assert executor.sessions == {}