diff --git a/README.md b/README.md index 00e33f8..a1463d2 100644 --- a/README.md +++ b/README.md @@ -122,8 +122,8 @@ The CLI passes both values to the Select AI SDK as `wallet_location` and The server accepts both A2A 1.x and the A2A v0.3 JSON-RPC streaming protocol for compatibility with Gemini Enterprise. -See the [A2A user guide](doc/source/user_guide/a2a.rst) for the dynamic -gateway, A2UI connection flow, persistent task state, task polling and +See the [A2A user guide](doc/source/user_guide/a2a.rst) for dynamic sessions, +the A2UI connection flow, persistent task state, task polling and cancellation, wallet configuration, and Google Cloud deployment modes. Generate the A2A v0.3 Agent Card to paste into Gemini Enterprise after the diff --git a/doc/source/image/a2a_architecture.svg b/doc/source/image/a2a_architecture.svg index c988b00..a9b9dfc 100644 --- a/doc/source/image/a2a_architecture.svg +++ b/doc/source/image/a2a_architecture.svg @@ -56,8 +56,8 @@ Oracle Database - SELECT_AI_A2A_TASKS - SELECT_AI_A2A_CONTEXTS + DBMS_AI_A2A_TASKS[$] + DBMS_AI_A2A_CONTEXTS[$] conversation history @@ -72,7 +72,7 @@ A2UI-capable client - select-ai a2a gateway + a2a serve --deployment clustered public Agent Card and A2UI A2A JSON-RPC, no streaming session/task proxy diff --git a/doc/source/user_guide/a2a.rst b/doc/source/user_guide/a2a.rst index f6a8eb5..f4a83ca 100644 --- a/doc/source/user_guide/a2a.rst +++ b/doc/source/user_guide/a2a.rst @@ -14,10 +14,10 @@ There are two deployment modes: * **Standalone A2A server**: one server owns a configured Oracle connection pool and one configured Select AI team. This mode supports streaming and is suitable when the service owner controls the database identity. -* **Dynamic A2A gateway**: a public gateway asks the client for a database DSN, - username, password, and team name through an A2UI form. It opens a temporary - isolated worker session for that selection. This mode supports task polling, - but does not advertise or implement streaming. +* **Clustered A2A server**: the public server asks the client for any database + DSN, username, password, and team values not fixed at deployment time. It + opens a temporary isolated worker session for that selection. This mode + supports task polling, but does not advertise or implement streaming. .. only:: html @@ -47,7 +47,8 @@ that can reach the required Oracle Database and, for dynamic deployments, the Consul service: * ``select-ai a2a serve`` runs the standalone server. -* ``select-ai a2a gateway`` runs the public dynamic gateway. +* ``select-ai a2a serve --deployment clustered`` runs the public dynamic + server. * ``select-ai a2a worker`` runs a database-bearing dynamic worker. The commands can be packaged into the platform's preferred container or @@ -77,7 +78,7 @@ The public routes are the same in both modes: - Selected by the client for each session. * - Public process - ``select-ai a2a serve`` - - ``select-ai a2a gateway`` + - ``select-ai a2a serve --deployment clustered`` * - Session process - The server's shared asynchronous connection pool. - A worker child process and connection pool per active session. @@ -180,6 +181,34 @@ The server prints the discovery URL when it starts: The exact Uvicorn startup lines vary by version and configuration. +End-user OAuth and session ownership +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +By default, ``select-ai a2a serve`` does not require an application OAuth +token. It gives each A2A conversation its own database session. A hosting +platform such as private Cloud Run can still authenticate the calling service, +but the A2A server does not verify which human user owns the conversation. + +Add ``--require-oauth`` when verified end-user ownership is required: + +.. code-block:: bash + + select-ai a2a serve \ + --require-oauth \ + --team ORACLE_AI_DATABASE_AGENT \ + --host 0.0.0.0 \ + --port 8000 \ + --public-url https://a2a.example.com + +Every A2A request must then contain ``Authorization: Bearer ...``. The server +uses both the authenticated owner and the A2A conversation to isolate +sessions, tasks, and conversation state. It accepts an OpenID Connect ID token +or an opaque OAuth 2.0 access token. ID-token ownership remains stable across +refreshes through the ``iss`` and ``sub`` claims; refreshing an opaque access +token starts a new owner and session scope. Tokens are never stored or logged. + +The Agent Card advertises bearer security only with ``--require-oauth``. + Agent Card discovery ~~~~~~~~~~~~~~~~~~~~ @@ -285,19 +314,23 @@ Persistent Oracle-backed state Both A2A deployment modes use Oracle implementations of the A2A task and context stores. In standalone mode, the server initializes these stores when -the application starts. In dynamic gateway mode, ``select-ai a2a gateway`` -does not connect to Oracle itself; the ``select-ai a2a worker`` command starts +the application starts. In clustered mode, +``select-ai a2a serve --deployment clustered`` does not connect to Oracle +itself; the ``select-ai a2a worker`` command starts the internal worker, and each connected worker session initializes the stores after it opens its supplied database connection. On first initialization, the stores create these tables if they do not already exist: -``SELECT_AI_A2A_TASKS`` +``DBMS_AI_A2A_TASKS$`` and ``DBMS_AI_A2A_TASKS`` Stores the task ID, context ID, serialized task JSON, owner, and update timestamp. It supports task retrieval, filtering, listing, pagination, and - deletion. + deletion. The object with the ``$`` suffix is the writable internal table; + the object without it is a read-only view of the same columns. -``SELECT_AI_A2A_CONTEXTS`` - Maps an A2A context and owner to an Oracle conversation ID. +``DBMS_AI_A2A_CONTEXTS$`` and ``DBMS_AI_A2A_CONTEXTS`` + Maps an A2A context and owner to an Oracle conversation ID. The object with + the ``$`` suffix is the writable internal table; the object without it is a + read-only view of the same columns. When a request starts a new context, Select AI creates an ``AsyncConversation`` and passes its ID to ``AsyncTeam.run``. Later messages in @@ -379,7 +412,7 @@ The ``-dev`` agent below is intentionally for local development only. It uses an in-memory, single-node Consul server bound to loopback. For the Google Cloud deployment, do not install or run a local agent: the deployment creates Consul in GKE from the -`Consul manifest `__. +`Consul manifest `__. Start the three local components in separate terminals: @@ -389,27 +422,33 @@ Start the three local components in separate terminals: consul agent -dev -bind=127.0.0.1 -client=127.0.0.1 # Terminal 2: worker - CONSUL_HTTP_URL=http://127.0.0.1:8500 \ - WORKER_ID=local-worker \ - WORKER_ADDRESS=127.0.0.1 \ - WORKER_PORT=8081 \ select-ai a2a worker \ --host 127.0.0.1 \ - --port 8081 + --port 8081 \ + --consul-url http://127.0.0.1:8500 \ + --worker-id local-worker \ + --worker-endpoint http://127.0.0.1:8081 # Terminal 3: public gateway - select-ai a2a gateway \ + select-ai a2a serve \ + --deployment clustered \ --host 127.0.0.1 \ --port 8000 \ - --agent-url http://127.0.0.1:8000 \ + --public-url http://127.0.0.1:8000 \ --consul-url http://127.0.0.1:8500 -The worker uses ``CONSUL_HTTP_URL`` (default ``http://consul:8500``), -``WORKER_ID``, ``WORKER_ADDRESS``, ``WORKER_PORT``, and the optional -``WORKER_ENDPOINT`` environment variables when registering with Consul. The -gateway uses ``AGENT_URL``, ``CONSUL_HTTP_URL``, ``WORKER_SERVICE`` (default -``select-ai-worker``), and ``SESSION_TTL_SECONDS``. The command-line options -override the corresponding environment variables. +The worker uses ``--consul-url``, ``--worker-id``, and the optional +``--worker-endpoint`` when registering with Consul. Their environment +fallbacks are ``CONSUL_HTTP_URL``, ``WORKER_ID``, and ``WORKER_ENDPOINT``; +explicit command-line values take precedence. ``--port`` is both the listening +port and the registered port. Cluster manifests can set ``WORKER_ADDRESS`` to +the pod IP when no advertised endpoint is supplied. Clustered serve uses +``PUBLIC_URL``, ``CONSUL_HTTP_URL``, ``WORKER_SERVICE`` (default +``select-ai-a2a-worker``), and ``SESSION_TTL_SECONDS``. + +Pass any fixed connection values with ``--dsn``, ``--user``, ``--password``, +and ``--team``. The generated form contains only the remaining values. Use +``--a2ui-form`` to supply a validated custom form for those missing values. The worker must be able to resolve the submitted DSN. For a TNS alias, set ``TNS_ADMIN`` in the worker terminal before starting it. The gateway session @@ -424,9 +463,10 @@ operations are rejected because the gateway does not proxy a live stream. A2UI database-connection form ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -The first ``message/send`` for a context returns a temporary task containing a -database connection form. The form is an A2UI v0.9 data artifact with the -``application/json+a2ui`` MIME type. It contains these fields: +When connection properties are missing, the first ``message/send`` for a +context returns a temporary task containing a database connection form. The +form is an A2UI v0.9 data artifact with the ``application/json+a2ui`` MIME +type. It contains only the values not configured on the server: .. list-table:: Dynamic connection form :header-rows: 1 @@ -435,17 +475,33 @@ database connection form. The form is an A2UI v0.9 data artifact with the * - Field - Purpose - * - ``dsn`` + * - ``connection_url`` - Oracle connect descriptor, simplified connect string, or TNS alias available to the worker. * - ``username`` - Database user for this session. * - ``password`` - Database password. The form renders this field as obscured input. - * - ``team_name`` + * - ``ai_agent`` - Select AI Agent Team to run in this session. * - ``submit_database_connection`` - - A2UI action that submits the four values to the gateway. + - A2UI action that submits exactly the displayed values. + +For example, a clustered server started with ``--dsn`` and ``--team`` renders +only username and password. Server-configured values are immutable and cannot +be overridden by a form submission. When all four values are configured, the +server skips the form and opens the worker session automatically. + +Use ``--a2ui-form PATH`` to replace the generated form. The file must be a +top-level JSON array of A2UI v0.9 operations using the advertised catalog and +one consistent template ``surfaceId``. It must contain exactly one +``submit_database_connection`` action whose context submits exactly the +missing canonical properties: ``connection_url``, ``username``, ``password``, +and/or ``ai_agent``. A collected password must use an obscured ``TextField`` +and the +template must not contain a password default. The server validates the file at +startup and replaces its template surface ID with a fresh ID every time the +form is displayed. After the action is submitted, the gateway selects a healthy worker through Consul and opens a session. The worker starts a child process, calls @@ -460,9 +516,9 @@ session routing metadata. The credentials remain in the worker child process for the lifetime of that session, so use TLS for client-to-gateway traffic and follow the security policies for any client that renders the form. -The gateway samples perform this handshake automatically. They use the +The dynamic-session samples perform this handshake automatically. They use the repository-local helper module -`samples/a2a/gateway/_common.py `__; +`samples/a2a/dynamic/_common.py `__; ``_common`` is not a package that users install with ``pip``. Running the samples from the repository root as shown below makes that helper available automatically. @@ -470,7 +526,7 @@ repository root as shown below makes that helper available automatically. The helper functions are: ``call(method, params)`` - Sends one A2A v0.3 JSON-RPC request to the configured gateway endpoint. + Sends one A2A v0.3 JSON-RPC request to the configured server endpoint. ``connect(prompt)`` Sends the initial prompt, reads the A2UI connection form, submits the @@ -479,26 +535,26 @@ The helper functions are: A2A context ID for the connected session. ``send_prompt(prompt, context_id, blocking=None)`` - Sends a database prompt in an existing gateway context. Passing + Sends a database prompt in an existing dynamic session. Passing ``blocking=False`` adds the non-blocking request option. ``print_task_summary(task)`` Prints the task state, result artifact name, and text parts without dumping the connection-form details. -If you copy a gateway sample into another directory, copy the -`_common.py helper `__ +If you copy a dynamic-session sample into another directory, copy the +`_common.py helper `__ with it, or replace these helpers with an A2A client implementation of your own. -The complete gateway setup is also documented in the -`gateway sample README `__. +The complete dynamic-session setup is also documented in the +`dynamic sample README `__. The examples below are the -`blocking_task.py sample `__ +`blocking_task.py sample `__ and the -`task_poll.py sample `__. +`task_poll.py sample `__. -.. literalinclude:: ../../../samples/a2a/gateway/blocking_task.py +.. literalinclude:: ../../../samples/a2a/dynamic/blocking_task.py :language: python :lines: 8- @@ -515,7 +571,7 @@ Run it with: .. code-block:: bash - python samples/a2a/gateway/blocking_task.py + python samples/a2a/dynamic/blocking_task.py Representative output is: @@ -525,10 +581,10 @@ Representative output is: Artifact: database-agent-result The database contains ... -The gateway polling sample uses the same connection-form handshake, then +The dynamic polling sample uses the same connection-form handshake, then passes ``configuration.blocking: false`` and polls ``tasks/get``: -.. literalinclude:: ../../../samples/a2a/gateway/task_poll.py +.. literalinclude:: ../../../samples/a2a/dynamic/task_poll.py :language: python :lines: 8- @@ -562,7 +618,8 @@ Consul is used for two kinds of routing metadata: * the worker service registration and health TTL, which let the gateway select a passing worker; and -* ``select-ai/sessions/`` and ``select-ai/tasks/`` key-value records, which +* ``select-ai-a2a/sessions/`` and ``select-ai-a2a/tasks/`` key-value records, + which route a context or task back to the worker that owns it. The task and conversation contents are not stored in Consul. They are stored @@ -579,8 +636,9 @@ All three gateway files are required together: .. code-block:: bash - select-ai a2a gateway \ - --agent-url https://gateway.example.com \ + select-ai a2a serve \ + --deployment clustered \ + --public-url https://gateway.example.com \ --worker-tls-ca-file /run/secrets/worker-ca.pem \ --worker-tls-cert-file /run/secrets/gateway-client.crt \ --worker-tls-key-file /run/secrets/gateway-client.key @@ -590,11 +648,12 @@ Start the worker with the matching server certificate, key, and CA: .. code-block:: bash select-ai a2a worker \ + --worker-endpoint https://worker.example.com:8443 \ --tls-cert-file /run/secrets/worker.crt \ --tls-key-file /run/secrets/worker.key \ --tls-ca-file /run/secrets/gateway-client-ca.pem -The worker must register an HTTPS endpoint through ``WORKER_ENDPOINT`` when +The worker must register an HTTPS endpoint through ``--worker-endpoint`` when mTLS is enabled. The gateway verifies the worker certificate with the CA and presents its client certificate. This mTLS option protects only the gateway-to-worker HTTP hop. It does not add wallet-based mTLS to the gateway's @@ -613,21 +672,25 @@ Standalone Google Cloud deployment ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ `gcloud/standalone/deploy.sh `__ -deploys one private Cloud Run service for one -database and team: +deploys one private Cloud Run service. Select its connection provisioning +explicitly: .. code-block:: bash gcloud/standalone/deploy.sh \ --project PROJECT_ID \ - --a2a-team ORACLE_AI_DATABASE_AGENT \ + --connection-mode dynamic \ --build The script creates or reuses an Artifact Registry repository, runtime service -account, Secret Manager secrets, and Cloud Run service. On the first run it -prompts for the database user, password, and connect descriptor. It injects -them into the service as ``SELECT_AI_USER``, ``SELECT_AI_PASSWORD``, and -``SELECT_AI_DB_CONNECT_STRING``. The service is private; the script grants +account, any requested Secret Manager secrets, and Cloud Run service. With no +database options, dynamic mode fixes neither the Connection URL nor AI Agent; +its A2UI form requests Connection URL, username, password, and AI Agent. Pass +``--db-dsn-secret`` or ``--a2a-team`` to fix either value independently, or +``--rotate-db-config`` to configure the deployment's Connection URL. Fixed +mode prompts for Connection URL, username, and password, requires +``--a2a-team``, and creates the shared startup pool. The service is private; the +script grants ``run.routes.invoke`` to the active deployment identity and the Gemini Enterprise Discovery Engine service agent. @@ -641,13 +704,20 @@ already deployed and only updates Cloud Run configuration or secrets. Important standalone options include: +* ``--connection-mode``: required ``fixed`` or ``dynamic`` provisioning. * ``--service``: Cloud Run service name; use a different service for each - fixed team/database deployment. -* ``--a2a-team``: team installed in Oracle Database. + deployment profile. +* ``--a2a-team``: fix the AI Agent in the deployment; required in fixed mode. +* ``--require-oauth``: require a Gemini Enterprise end-user OAuth bearer token + and scope state by verified owner as well as A2A conversation. * ``--pool-max-size``: maximum Oracle connections per Cloud Run instance. -* ``--max-instances``: Cloud Run instance limit. -* ``--wallet-archive``: upload or replace an Autonomous Database wallet ZIP. -* ``--rotate-db-credentials``: prompt for and rotate the database secrets. +* ``--max-instances``: Cloud Run instance limit; dynamic mode requires one. +* ``--wallet-archive``: upload or replace an Autonomous Database wallet ZIP + in fixed mode. +* ``--db-dsn-secret``: fix the Connection URL using an existing Secret Manager + secret. +* ``--rotate-db-config``: prompt for and create or rotate the deployment's + database values. For a wallet deployment, ``--wallet-archive`` stores the ZIP and wallet password in service-specific Secret Manager secrets. Cloud Run mounts the ZIP @@ -659,42 +729,50 @@ The standalone script prints the deployed Agent Card JSON at the end. The service is ready for a client when the printed card points to the final Cloud Run URL and the client can invoke the private service. -Dynamic gateway Google Cloud deployment +The Google Cloud scripts do not require end-user OAuth by default. Cloud Run +IAM authenticates Gemini Enterprise, and ``a2a serve`` separates database +sessions by A2A conversation without asserting a verified human owner. Add +``--require-oauth`` to the deployment command when verified user ownership is +required. Gemini Enterprise must then be configured with end-user OAuth rather +than **Skip & Finish**. It uses ``X-Serverless-Authorization`` for Cloud Run +IAM and sends the user token separately in ``Authorization``. + +Clustered Google Cloud deployment ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -`gcloud/gateway/deploy.sh `__ -deploys the gateway topology: Cloud Run for the -public gateway, GKE Autopilot for worker replicas, and Consul for discovery and +`gcloud/cluster/deploy.sh `__ +deploys the clustered topology: Cloud Run for the +public server, GKE Autopilot for worker replicas, and Consul for discovery and routing. .. code-block:: bash - gcloud/gateway/deploy.sh \ + gcloud/cluster/deploy.sh \ --project PROJECT_ID \ --region us-central1 \ --worker-replicas 3 The script and -`gcloud/gateway/cloudbuild.yaml `__ +`gcloud/cluster/cloudbuild.yaml `__ perform these steps: 1. Enable the required Google Cloud APIs and create or reuse Artifact Registry and a GKE Autopilot cluster. 2. Configure additive VPC DNS so StatefulSet worker names resolve to current Pod IPs after a worker is recreated. -3. Deploy the gateway namespace, internal Consul service, and worker service. +3. Deploy the A2A namespace, internal Consul service, and worker service. 4. Build and publish the image with Cloud Build. 5. Deploy the requested worker replicas with ``select-ai a2a worker`` and the selected session TTL. -6. Deploy Cloud Run with ``select-ai a2a gateway`` and set ``AGENT_URL`` to the - final Cloud Run URL. +6. Deploy Cloud Run with ``select-ai a2a serve --deployment clustered`` and + set ``PUBLIC_URL`` to the final Cloud Run URL. -The Cloud Run gateway uses direct VPC egress to reach the internal Consul +The Cloud Run server uses direct VPC egress to reach the internal Consul load-balancer address and worker endpoints. The GKE worker service is headless service discovery, not a public load balancer. Consul selects a healthy worker when a session opens; the route is then pinned to that worker. -Important gateway options include: +Important clustered options include: * ``--cluster`` and ``--gke-dns-domain``: GKE Autopilot cluster and immutable additive DNS domain. @@ -702,18 +780,18 @@ Important gateway options include: internal worker infrastructure. * ``--worker-replicas``: number of worker runtimes available for new sessions. * ``--session-ttl-seconds``: lifetime of temporary database sessions. -* ``--enable-worker-mtls``: create and use gateway-to-worker certificates. +* ``--enable-worker-mtls``: create and use server-to-worker certificates. * ``--rotate-worker-mtls``: replace the existing test certificates and restart the worker workload. -The default gateway deployment uses private-VPC HTTP between Cloud Run and +The default clustered deployment uses private-VPC HTTP between Cloud Run and GKE. With ``--enable-worker-mtls``, the script creates short-lived test PKI material in Secret Manager, creates the Kubernetes TLS secrets, deploys the -worker StatefulSet variant, and mounts the gateway client certificate into +worker StatefulSet variant, and mounts the server client certificate into Cloud Run. The worker certificate uses the GKE StatefulSet DNS name, so the ``--gke-dns-domain`` value must remain consistent with the cluster. -The gateway deployment does not upload an Oracle wallet because the dynamic +The clustered deployment does not upload an Oracle wallet because the dynamic worker session path currently accepts only DSN/user/password. If database mTLS is required, use the standalone deployment or provide a separate database connection mechanism to the worker implementation. @@ -733,23 +811,23 @@ The deployment files are intended to be read together: - Builds and publishes the standalone container image. * - `gcloud/standalone/README.md `__ - Documents IAM, secrets, wallet archives, and update behavior. - * - `gcloud/gateway/deploy.sh `__ - - Creates or reuses the gateway GKE/Cloud Run topology and supplies build + * - `gcloud/cluster/deploy.sh `__ + - Creates or reuses the clustered GKE/Cloud Run topology and supplies build substitutions. - * - `gcloud/gateway/cloudbuild.yaml `__ + * - `gcloud/cluster/cloudbuild.yaml `__ - Builds the image, deploys Consul and workers, and deploys Cloud Run. - * - `gcloud/gateway/gke manifests `__ + * - `gcloud/cluster/gke manifests `__ - Namespace, Consul, headless worker service, and HTTP or mTLS worker workloads. - * - `gcloud/gateway/README.md `__ + * - `gcloud/cluster/README.md `__ - Explains the topology, DNS, mTLS test mode, and operational details. Troubleshooting and security notes ================================== -* If the Agent Card advertises the wrong URL, set ``--public-url`` for the - standalone server or ``AGENT_URL`` for the gateway. The URL must be the - client-visible base URL, not an internal container address. +* If the Agent Card advertises the wrong URL, set ``--public-url`` or + ``PUBLIC_URL``. The URL must be the client-visible base URL, not an internal + container address. * If the gateway returns the connection form repeatedly, check that the worker is passing in Consul, that the worker can resolve and connect to Oracle, and that the session TTL has not expired. diff --git a/doc/source/user_guide/cli.rst b/doc/source/user_guide/cli.rst index 223675e..84abdb1 100644 --- a/doc/source/user_guide/cli.rst +++ b/doc/source/user_guide/cli.rst @@ -206,7 +206,6 @@ current defaults: select-ai a2a --help select-ai a2a serve --help - select-ai a2a gateway --help select-ai a2a worker --help select-ai a2a agent-card --help @@ -218,11 +217,9 @@ current defaults: * - Command - Purpose * - ``select-ai a2a serve`` - - Start a standalone A2A HTTP server for one configured database AI Agent - Team. It accepts the database connection and optional wallet options. - * - ``select-ai a2a gateway`` - - Start the public dynamic A2A/A2UI gateway. It uses Consul to discover - workers and does not connect to Oracle directly. + - Start the public A2A HTTP server in standalone or clustered deployment + mode. Standalone accepts database connection and wallet options; + clustered uses Consul to discover workers. * - ``select-ai a2a worker`` - Start the internal worker that registers with Consul and creates an isolated database-bearing child process for each submitted connection. @@ -260,6 +257,12 @@ Important options are ``--team`` (required), ``--host``, ``--port``, wallet for this standalone path. If no password is provided, the command prompts for it. +By default, ``a2a serve`` does not require an application OAuth token and +separates database sessions by A2A conversation. Add ``--require-oauth`` to +require ``Authorization: Bearer ...`` and scope sessions by authenticated +owner as well as conversation. The Agent Card advertises bearer security only +in that mode. + Dynamic gateway and worker -------------------------- @@ -270,29 +273,38 @@ passing the A2UI form values to the worker. .. code-block:: bash - CONSUL_HTTP_URL=http://127.0.0.1:8500 \ - WORKER_ID=local-worker \ - WORKER_ADDRESS=127.0.0.1 \ - WORKER_PORT=8081 \ - select-ai a2a worker --host 127.0.0.1 --port 8081 + select-ai a2a worker \ + --host 127.0.0.1 \ + --port 8081 \ + --consul-url http://127.0.0.1:8500 \ + --worker-id local-worker \ + --worker-endpoint http://127.0.0.1:8081 - select-ai a2a gateway \ + select-ai a2a serve \ + --deployment clustered \ --host 127.0.0.1 \ --port 8000 \ - --agent-url http://127.0.0.1:8000 \ + --public-url http://127.0.0.1:8000 \ --consul-url http://127.0.0.1:8500 -The worker options are ``--host``, ``--port``, ``--session-ttl-seconds``, and -``--session-start-timeout-seconds``. Its registration can be configured with -the ``CONSUL_HTTP_URL``, ``WORKER_ID``, ``WORKER_ADDRESS``, ``WORKER_PORT``, -and optional ``WORKER_ENDPOINT`` environment variables. The worker's -``--tls-cert-file``, ``--tls-key-file``, and ``--tls-ca-file`` options enable -gateway-to-worker mTLS; provide all three together. - -The gateway options are ``--agent-url`` (required, or ``AGENT_URL``), +The worker options are ``--host``, ``--port``, ``--worker-id``, +``--consul-url``, ``--worker-endpoint``, ``--session-ttl-seconds``, and +``--session-start-timeout-seconds``. ``WORKER_ID``, ``CONSUL_HTTP_URL``, and +``WORKER_ENDPOINT`` are environment fallbacks; explicit command-line values +take precedence. ``--port`` is also the port registered with Consul. Cluster +manifests can set ``WORKER_ADDRESS`` to the pod IP when no endpoint is +supplied. The worker's ``--tls-cert-file``, ``--tls-key-file``, and +``--tls-ca-file`` options enable gateway-to-worker mTLS; provide all three +together. + +The clustered serve options are ``--public-url`` (or ``PUBLIC_URL``), ``--consul-url`` (or ``CONSUL_HTTP_URL``), ``--worker-service`` (or ``WORKER_SERVICE``), and ``--session-ttl-seconds`` (or -``SESSION_TTL_SECONDS``). The optional +``SESSION_TTL_SECONDS``). ``--dsn``, ``--user``, ``--password``, and ``--team`` +fix any supplied connection values; the generated A2UI form contains only the +missing values. ``--a2ui-form`` loads a custom A2UI JSON form that must submit +exactly those missing values. ``--require-oauth`` has the same ownership +behavior in standalone and clustered deployments. The optional ``--worker-tls-ca-file``, ``--worker-tls-cert-file``, and ``--worker-tls-key-file`` options configure the gateway's client side of gateway-to-worker mTLS. These TLS settings protect the internal HTTP hop and @@ -341,9 +353,7 @@ Command summary * - ``select-ai profile translate`` - Translate text with a saved profile. * - ``select-ai a2a serve`` - - Start the standalone A2A server. - * - ``select-ai a2a gateway`` - - Start the public dynamic A2A gateway. + - Start the public A2A server in standalone or clustered mode. * - ``select-ai a2a worker`` - Start the internal dynamic-session worker. * - ``select-ai a2a agent-card`` diff --git a/docker/a2a-entrypoint.sh b/docker/a2a-entrypoint.sh index 258a9b1..8f241a8 100644 --- a/docker/a2a-entrypoint.sh +++ b/docker/a2a-entrypoint.sh @@ -28,13 +28,21 @@ if [ -f "$wallet_archive" ]; then export SELECT_AI_WALLET_LOCATION="$(dirname "$wallet_file")" fi -: "${SELECT_AI_A2A_TEAM:?SELECT_AI_A2A_TEAM is required}" : "${PUBLIC_URL:?PUBLIC_URL is required}" : "${SELECT_AI_POOL_MAX_SIZE:=10}" -exec select-ai a2a serve \ - --team "$SELECT_AI_A2A_TEAM" \ - --host 0.0.0.0 \ - --port "${PORT:-8080}" \ - --pool-max-size "$SELECT_AI_POOL_MAX_SIZE" \ +serve_args=( + --deployment standalone + --host 0.0.0.0 + --port "${PORT:-8080}" + --pool-max-size "$SELECT_AI_POOL_MAX_SIZE" --public-url "$PUBLIC_URL" +) +if [ -n "${SELECT_AI_A2A_TEAM:-}" ]; then + serve_args+=(--team "$SELECT_AI_A2A_TEAM") +fi +if [ "${SELECT_AI_A2A_REQUIRE_OAUTH:-false}" = "true" ]; then + serve_args+=(--require-oauth) +fi + +exec select-ai a2a serve "${serve_args[@]}" diff --git a/gcloud/README.md b/gcloud/README.md index bf005a8..01ae2bd 100644 --- a/gcloud/README.md +++ b/gcloud/README.md @@ -1,115 +1,107 @@ # Google Cloud deployment modes -Select AI for Python supports two A2A deployment architectures. The key -decisions are where the database connection and Select AI team are selected, -which components carry the session, and how capacity is added: - -- Standalone fixes the database and team at deployment time. One Cloud Run A2A - service owns the configured connection pool and serves that team. -- The gateway selects the database and team per user session. A Cloud Run A2A - gateway routes sessions through Consul to a clustered GKE worker pool, with - one isolated child runtime and database connection pool per active session. - -The gateway architecture is designed for horizontal session capacity. Gateway -instances, Consul, and worker replicas are separate components; adding worker -replicas increases the number of concurrent database sessions that can be -hosted behind the same A2A endpoint. Consul preserves session and task affinity -when requests reach different gateway instances. Oracle Database capacity and -the configured session TTL remain the limiting factors. +Select AI for Python supports three Google Cloud A2A deployment profiles. The +profiles keep deployment topology separate from connection provisioning: + +1. **Fixed standalone** runs one Cloud Run service with DSN, username, + password, and team fixed by deployment. It uses the shared connection pool + and supports streaming. +2. **Dynamic standalone** runs one Cloud Run service. Connection URL and AI + Agent may be fixed independently by deployment; each user supplies all + remaining connection values through A2UI and receives an isolated + child-process session. Its in-memory routing requires exactly one Cloud Run + instance and streaming is disabled. +3. **Clustered dynamic** runs the same public A2A contract on Cloud Run, with + Consul and workers on GKE. DSN and team are fixed; each user supplies + username and password through A2UI. Consul provides distributed routing and + workers own the isolated session processes. ![Select AI A2A deployment architecture](../doc/source/image/a2a_architecture.svg) -The A2A commands are cloud-neutral: `select-ai a2a serve`, -`select-ai a2a gateway`, and `select-ai a2a worker` can run as processes or -containers on any cloud platform, a Kubernetes cluster, or self-managed -infrastructure with the required Oracle and Consul connectivity. The scripts -in this directory are optional Google Cloud automation for the Cloud Run/GKE -topologies shown below. +The A2A commands are cloud-neutral. These scripts are optional Google Cloud +automation for the Cloud Run and GKE profiles. -## What the A2A client connects to +## Standalone deployments -### Standalone server +Use the same script with an explicit connection mode. -The standalone deployment is one Cloud Run A2A service for one configured -Oracle database and one Select AI team. - -The service receives its database credentials from Secret Manager. The A2A -client can discover the Agent Card and immediately send a database prompt. -The server supports blocking tasks, task polling, and streaming responses. - -Deploy it with: +Dynamic standalone: ```bash -gcloud/standalone/deploy.sh --build +gcloud/standalone/deploy.sh \ + --project PROJECT_ID \ + --connection-mode dynamic \ + --build ``` -Use [standalone deployment](standalone/README.md) for the deployment details. +Fixed standalone: -### Dynamic gateway +```bash +gcloud/standalone/deploy.sh \ + --project PROJECT_ID \ + --connection-mode fixed \ + --a2a-team ORACLE_AI_DATABASE_AGENT \ + --build +``` -The gateway deployment provides one public A2A endpoint for users who choose -the database and Select AI team at runtime. +The modes use distinct default services: -The client first sends a message and receives an A2UI connection form. After -the client submits the DSN, username, password, and team name, the gateway -opens a temporary worker session. Subsequent A2A messages use that session and -execute against the selected database and team. +- `select-ai-a2a-standalone-dynamic` +- `select-ai-a2a-standalone-fixed` -The gateway supports blocking tasks and asynchronous task polling. Its Agent -Card advertises `streaming: false`; clients use `message/send` followed by -`tasks/get` for long-running work. The gateway-to-worker path uses internal -protobuf messages, while the public client-facing path remains A2A JSON-RPC. +See the [standalone deployment guide](standalone/README.md). -The gateway database session currently accepts a DSN, username, and password. -Wallet-based Oracle Database mTLS is not yet supported by this session path. -The optional mTLS deployment mode described in the gateway documentation -secures the gateway-to-worker connection; it is separate from database mTLS. +The standalone guide also shows how to select the newest existing image +directly from Artifact Registry by immutable digest. Use that workflow when +creating another Cloud Run service without rebuilding an identical image or +depending on an existing service as the image source. -Deploy it with: +## Clustered dynamic deployment ```bash -gcloud/gateway/deploy.sh --project PROJECT_ID +gcloud/cluster/deploy.sh --project PROJECT_ID ``` -Use [gateway deployment](gateway/README.md) for the deployment details. +Pass `--image-uri IMAGE@sha256:DIGEST` to reuse an existing immutable image +for both the Cloud Run server and GKE workers instead of building another one. + +The default GKE cluster and public Cloud Run service are both named +`select-ai-a2a-cluster` in their respective resource namespaces. The server uses direct VPC egress to +reach Consul and the GKE workers. The worker transport is private HTTP by +default; optional mTLS protects the server-to-worker hop. Wallet-based Oracle +Database mTLS is not yet supported for dynamic sessions. + +See the [cluster deployment guide](cluster/README.md). ## Client-visible differences -| Client concern | Standalone server | Dynamic gateway | -| --- | --- | --- | -| Database/team selection | Configured by the deployment | Submitted by each user session through A2UI | -| First client operation | Send the database prompt | Send a prompt, submit the connection form, then send the database prompt | -| Credentials | Stored in Secret Manager for the service | Supplied for the temporary session and held by its worker | -| Database mTLS | Supported through the standalone wallet configuration | Not yet supported for gateway database sessions | -| Public service | One Cloud Run A2A service | Cloud Run gateway backed by GKE workers and Consul | -| Agent Card input | `text/plain` | `text/plain` and `application/json+a2ui` | -| Agent Card streaming | `true` | `false` | -| Blocking request | `message/send` waits for the final task result | `message/send` waits for the final task result after the session is connected | -| Streaming response | Supported through A2A streaming methods and SSE | Not available; clients use task polling | -| Asynchronous task | `message/send` with `configuration.blocking: false` | `message/send` with `configuration.blocking: false` | -| Task polling | `tasks/get` until the task reaches a terminal state | `tasks/get` until the task reaches a terminal state | -| Session ownership | Cloud Run service database pool | One child process and async pool per active user session | -| Task/context storage | Oracle Database | Oracle Database, with Consul routing metadata | -| Capacity control | Cloud Run instances and per-instance pool size | Gateway instances, Consul routing, worker replicas, per-session pools, and session TTL | -| Best fit | One known database/team and predictable operations | Multiple databases/teams selected dynamically from one endpoint | - -Both deployments expose the public A2A endpoint at: +| Concern | Fixed standalone | Dynamic standalone | Clustered dynamic | +| --- | --- | --- | --- | +| Topology | Cloud Run | One Cloud Run instance | Cloud Run + Consul + GKE workers | +| Deployment-fixed values | Connection URL, username, password, AI Agent | Any subset of Connection URL and AI Agent | Connection URL and AI Agent | +| A2UI form | None | All connection values not fixed by deployment | Username and password | +| Database runtime | Shared configured pool | Isolated local child process | Isolated worker child process | +| Streaming | Supported | Not currently supported | Not currently supported | +| Task/context storage | Oracle Database | Oracle Database | Oracle Database with Consul routing metadata | +| Scaling | Cloud Run instances and pool size | One Cloud Run instance | Cloud Run instances and worker replicas | + +All three expose: ```text /.well-known/agent-card.json /a2a/jsonrpc/ ``` -Both accept A2A 1.0 method names and the A2A v0.3 compatibility method names. -The gateway client flow is documented in the -[gateway samples](../samples/a2a/gateway/README.md). - -## Which deployment should you choose? +They accept A2A 1.0 method names and the A2A v0.3 compatibility method names. +The connection-form client flow is documented in the +[dynamic-session samples](../samples/a2a/dynamic/README.md). -Choose the standalone server when the service owner controls the database -identity and team, wants clients to send prompts immediately, and benefits -from streaming responses. +By default, the private Cloud Run deployments rely on Cloud Run IAM and keep +database sessions separate by A2A conversation. Add `--require-oauth` to +either deployment script when verified human-user ownership is required. In +that mode Gemini Enterprise must be configured with end-user OAuth and sends +its user token in `Authorization`; missing tokens receive HTTP 401. -Choose the gateway when one A2A endpoint must serve users selecting different -Oracle databases or teams, with isolated temporary sessions and worker-based -capacity. +`X-Serverless-Authorization` is separate and automatic: Gemini Enterprise uses +it to invoke the private Cloud Run service, and Cloud Run consumes it before +the request reaches `a2a serve`. diff --git a/gcloud/gateway/README.md b/gcloud/cluster/README.md similarity index 65% rename from gcloud/gateway/README.md rename to gcloud/cluster/README.md index 5e11273..44b1654 100644 --- a/gcloud/gateway/README.md +++ b/gcloud/cluster/README.md @@ -1,8 +1,9 @@ -# Dynamic gateway +# Clustered Select AI A2A deployment -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 +Clustered mode exposes one public A2A server. The deployment fixes the Oracle +database DSN and Select AI team; each user supplies a database username and +password through the generated A2UI 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. @@ -10,7 +11,7 @@ different lifetime in seconds with `--session-ttl-seconds`; for example, ```text ┌──────────────┐ A2A JSON-RPC/HTTP ┌──────────────────┐ ┌────────────────────────────┐ -│ A2A client │────────────────────►│ Gateway instances│──── route lookup/update ──────────► │ Service Registry │ +│ A2A client │────────────────────►│ Server instances │──── route lookup/update ──────────► │ Service Registry │ └──────────────┘ │ Public A2A API │ │ Service discovery │ │ A2UI bootstrap │ │ Session routes │ └────────┬─────────┘ │ Task routes │ @@ -36,7 +37,7 @@ different lifetime in seconds with `--session-ttl-seconds`; for example, └──────────────────┘ ``` -The gateway is the only public A2A application. It selects a worker through +The Cloud Run server is the only public A2A application. It selects a worker through the Service Registry, opens a session there, and proxies subsequent A2A calls using the internal protobuf protocol. The selected worker starts one child runtime for that session. The child owns the database connection, @@ -44,13 +45,13 @@ runtime for that session. The child owns the database connection, The session connection path currently accepts a DSN, username, and password. Oracle Database wallet-based mTLS is not yet supported for these dynamic -sessions. The optional worker mTLS mode below protects the gateway-to-worker +sessions. The optional worker mTLS mode below protects the server-to-worker HTTP connection; it does not provide database mTLS. The Service Registry stores only service-discovery and non-secret session/task-to-worker metadata. Task payloads and context mappings remain in Oracle. Connection-form tasks are response-only bootstrap tasks: they are -created by the gateway before a database session exists and are not persisted +created by the server before a database session exists and are not persisted or routed. ## GCP deployment @@ -64,7 +65,7 @@ The protocol architecture above is implemented on GCP as follows: │ public A2A v ┌────────────────────────────┐ - │ Cloud Run gateway │ + │ Cloud Run A2A server │ │ A2A proxy + form bootstrap │ └──────┬───────────┬─────────┘ │ │ private VPC: mTLS request to worker hostname @@ -88,61 +89,114 @@ The protocol architecture above is implemented on GCP as follows: Run this from the repository root: ```bash -gcloud/gateway/deploy.sh --project PROJECT_ID +gcloud/cluster/deploy.sh --project PROJECT_ID ``` +That command builds a new image. To reuse the newest existing Select AI image +from Artifact Registry, resolve its immutable digest URI first: + +```bash +PROJECT_ID=PROJECT_ID +REGION=us-central1 +REPOSITORY=select-ai +IMAGE_NAME=select-ai + +IMAGE_RECORD="$(gcloud artifacts docker images list \ + "$REGION-docker.pkg.dev/$PROJECT_ID/$REPOSITORY/$IMAGE_NAME" \ + --project "$PROJECT_ID" \ + --include-tags \ + --sort-by='~UPDATE_TIME' \ + --limit=1 \ + --format='csv[no-heading](package,version)')" + +if [[ -z "$IMAGE_RECORD" ]]; then + echo "No Select AI image exists in Artifact Registry." >&2 + exit 1 +fi + +IMAGE_URI="${IMAGE_RECORD/,/@}" + +gcloud/cluster/deploy.sh \ + --project "$PROJECT_ID" \ + --region "$REGION" \ + --repository "$REPOSITORY" \ + --a2a-team ORACLE_AI_DATABASE_AGENT \ + --db-dsn-secret select-ai-a2a-cluster-db-connect-string \ + --image-uri "$IMAGE_URI" +``` + +The first deployment prompts for the Connection URL, creates the GKE Autopilot +cluster and service-specific secret, and uses the same pinned image for the +GKE workers and Cloud Run server. Add `--require-oauth` to the final command +only when Gemini Enterprise end-user OAuth is configured for this clustered +agent. + +By default, Cloud Run IAM authenticates Gemini Enterprise and database +sessions are separated by A2A conversation; the server does not verify the +human user who owns a conversation. Add `--require-oauth` when the deployment +requires authenticated end-user ownership. Gemini Enterprise must then be +configured with end-user OAuth and every request must carry its separate +`Authorization: Bearer ...` token. Requests without it receive HTTP 401. + +On the first deployment, the script prompts for the fixed database DSN and +stores it in Secret Manager. Override the default secret name with +`--db-dsn-secret` to reuse an existing DSN secret. The team defaults to +`ORACLE_AI_DATABASE_AGENT` and can be changed with `--a2a-team`. Because DSN +and team are fixed by deployment, the generated A2UI form asks each user for +only username and password. + The script creates the Artifact Registry repository and GKE Autopilot cluster 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; +1. builds and publishes `docker/Dockerfile`, or reuses `--image-uri`; 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. +4. deploys the same image to Cloud Run using + `select-ai a2a serve --deployment clustered`; +5. sets the final Cloud Run URL in `PUBLIC_URL` for the Agent Card. Common options: ```bash -gcloud/gateway/deploy.sh \ +gcloud/cluster/deploy.sh \ --project PROJECT_ID \ --region us-central1 \ - --cluster select-ai-a2a-gateway \ - --gke-dns-domain select-ai-a2a-gateway.internal \ + --cluster select-ai-a2a-cluster \ + --gke-dns-domain select-ai-a2a-cluster.internal \ --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 gateway keeps only the connection +The Cloud Run server uses direct VPC egress to reach the internal Consul load +balancer and GKE worker pod addresses. The server keeps only the connection form task transiently, before a database session exists. Connected task and context state is stored in Oracle on the selected worker. Workers, rather -than the gateway, provide the clustered capacity for dynamic sessions. -The deployment currently keeps one gateway instance as an operational default; -the gateway does not cache forms or connected task/context state. Gateway +than the public server, provide the clustered capacity for dynamic sessions. +The public server does not cache forms or connected task/context state. Server scaling does not change session affinity because Consul stores the session and task routes. `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. +the selected image and Consul endpoint values to both the GKE workers and the +Cloud Run server 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. -This mTLS mode applies only between the Cloud Run gateway and GKE workers. It +This mTLS mode applies only between the Cloud Run server and GKE workers. It is independent of Oracle Database authentication, and does not enable wallet- -based database mTLS for gateway sessions. +based database mTLS for dynamic sessions. For a short-lived GCloud mTLS test: ```bash -gcloud/gateway/deploy.sh \ +gcloud/cluster/deploy.sh \ --project PROJECT_ID \ --enable-worker-mtls \ --mtls-cert-validity-days 365 @@ -156,7 +210,7 @@ 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. +certificates: the server 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 @@ -164,10 +218,10 @@ 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. The `select-ai-worker` headless Service gives +Workers run as a StatefulSet. The `select-ai-a2a-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 +`select-ai-a2a-worker-0.select-ai-a2a-worker.select-ai-a2a.svc.select-ai-a2a-cluster.internal`. +Consul registers that name, so the server 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. @@ -192,12 +246,13 @@ GKE gives a StatefulSet Pod a DNS name using this form: For this deployment, worker 0 is: ```text -select-ai-worker-0.select-ai-worker.select-ai-gateway.svc.select-ai-a2a-gateway.internal +select-ai-a2a-worker-0.select-ai-a2a-worker.select-ai-a2a.svc.select-ai-a2a-cluster.internal ``` -`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 +`select-ai-a2a-worker-0` is the StatefulSet Pod name, +`select-ai-a2a-worker` is the headless Service, `select-ai-a2a` is the +Kubernetes namespace, and `select-ai-a2a-cluster.internal` is the +`--gke-dns-domain` value. GKE updates the resulting record when the Pod IP changes. ### Certificate mounts @@ -206,11 +261,11 @@ Worker certificate files are mounted from Kubernetes Secrets: | 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 | +| `/var/run/select-ai-mtls/tls.crt` | `select-ai-a2a-worker-server-tls` | `tls.crt` | worker HTTPS server certificate | +| `/var/run/select-ai-mtls/tls.key` | `select-ai-a2a-worker-server-tls` | `tls.key` | worker HTTPS private key | +| `/var/run/select-ai-mtls/server-ca.crt` | `select-ai-a2a-server-client-ca` | `ca.crt` | validates the server client certificate | -The Cloud Run gateway certificate files are mounted from Google Secret Manager. +The Cloud Run server certificate files are mounted from Google Secret Manager. The identity that submits Cloud Build needs permission to use GKE, Cloud Run, and Secret Manager. GKE maintains the managed worker DNS records. diff --git a/gcloud/gateway/cloudbuild.yaml b/gcloud/cluster/cloudbuild.yaml similarity index 56% rename from gcloud/gateway/cloudbuild.yaml rename to gcloud/cluster/cloudbuild.yaml index 89bb0af..ad8d9aa 100644 --- a/gcloud/gateway/cloudbuild.yaml +++ b/gcloud/cluster/cloudbuild.yaml @@ -1,27 +1,25 @@ -# Build one Select AI image, deploy Consul and workers to GKE, then deploy the -# public dynamic gateway to Cloud Run. +# Build or reuse one Select AI image, deploy Consul and workers to GKE, then +# deploy the public clustered server to Cloud Run. steps: - name: gcr.io/cloud-builders/docker + entrypoint: bash 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} + - -ceu + - | + if [[ -n "${_IMAGE_URI}" ]]; then + echo "Reusing existing image: ${_IMAGE_URI}" + exit 0 + fi + image="${_REGION}-docker.pkg.dev/$PROJECT_ID/${_REPOSITORY}/select-ai:${_IMAGE_TAG}" + docker build --file docker/Dockerfile --tag "$image" . + # GKE and Cloud Run consume the image in later steps, so publish it + # before deploying either workload. + docker push "$image" - name: gcr.io/cloud-builders/gke-deploy args: - run - - --filename=gcloud/gateway/gke/namespace.yaml + - --filename=gcloud/cluster/gke/namespace.yaml - --location=${_REGION} - --cluster=${_CLUSTER} @@ -36,7 +34,7 @@ steps: - name: gcr.io/cloud-builders/gke-deploy args: - run - - --filename=gcloud/gateway/gke/worker-service.yaml + - --filename=gcloud/cluster/gke/worker-service.yaml - --location=${_REGION} - --cluster=${_CLUSTER} @@ -57,12 +55,12 @@ steps: - -ceu - | if [[ "${_ENABLE_WORKER_MTLS}" == "true" ]]; then - kubectl delete --namespace select-ai-gateway deployment/select-ai-worker --ignore-not-found + kubectl delete --namespace select-ai-a2a deployment/select-ai-a2a-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 + kubectl delete --namespace select-ai-a2a statefulset/select-ai-a2a-worker --ignore-not-found --wait=true fi else - kubectl delete --namespace select-ai-gateway statefulset/select-ai-worker --ignore-not-found + kubectl delete --namespace select-ai-a2a statefulset/select-ai-a2a-worker --ignore-not-found fi - name: gcr.io/cloud-builders/gcloud @@ -73,9 +71,9 @@ steps: 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 + gcloud secrets versions access latest --secret=select-ai-a2a-worker-mtls-cert > /workspace/worker.crt + gcloud secrets versions access latest --secret=select-ai-a2a-worker-mtls-key > /workspace/worker.key + gcloud secrets versions access latest --secret=select-ai-a2a-mtls-ca > /workspace/server-ca.crt - name: gcr.io/cloud-builders/kubectl entrypoint: bash @@ -88,17 +86,17 @@ steps: if [[ "${_ENABLE_WORKER_MTLS}" != "true" ]]; then exit 0 fi - kubectl -n select-ai-gateway create secret tls select-ai-worker-server-tls \ + kubectl -n select-ai-a2a create secret tls select-ai-a2a-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 \ + kubectl -n select-ai-a2a create secret generic select-ai-a2a-server-client-ca \ + --from-file=ca.crt=/workspace/server-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 + - --filename=gcloud/cluster/gke/consul.yaml - --location=${_REGION} - --cluster=${_CLUSTER} @@ -110,10 +108,10 @@ steps: args: - -ceu - | - kubectl wait --namespace select-ai-gateway \ + kubectl wait --namespace select-ai-a2a \ --for=jsonpath='{.status.loadBalancer.ingress[0].ip}' \ service/consul-server --timeout=300s - kubectl get --namespace select-ai-gateway service/consul-server \ + kubectl get --namespace select-ai-a2a service/consul-server \ -o jsonpath='{.status.loadBalancer.ingress[0].ip}' \ > /workspace/consul-ip @@ -122,10 +120,13 @@ steps: args: - -ceu - | - image="${_REGION}-docker.pkg.dev/$PROJECT_ID/${_REPOSITORY}/select-ai:${_IMAGE_TAG}" - worker_manifest=gcloud/gateway/gke/worker.yaml + image="${_IMAGE_URI}" + if [[ -z "$image" ]]; then + image="${_REGION}-docker.pkg.dev/$PROJECT_ID/${_REPOSITORY}/select-ai:${_IMAGE_TAG}" + fi + worker_manifest=gcloud/cluster/gke/worker.yaml if [[ "${_ENABLE_WORKER_MTLS}" == "true" ]]; then - worker_manifest=gcloud/gateway/gke/worker-mtls.yaml + worker_manifest=gcloud/cluster/gke/worker-mtls.yaml fi sed -e "s|WORKER_IMAGE|$image|" \ -e "s|WORKER_REPLICAS|${_WORKER_REPLICAS}|" \ @@ -155,9 +156,9 @@ steps: - -ceu - | if [[ "${_ENABLE_WORKER_MTLS}" == "true" ]]; then - kubectl rollout status --namespace select-ai-gateway statefulset/select-ai-worker --timeout=300s + kubectl rollout status --namespace select-ai-a2a statefulset/select-ai-a2a-worker --timeout=300s else - kubectl rollout status --namespace select-ai-gateway deployment/select-ai-worker --timeout=300s + kubectl rollout status --namespace select-ai-a2a deployment/select-ai-a2a-worker --timeout=300s fi - name: gcr.io/cloud-builders/gcloud @@ -166,51 +167,64 @@ steps: - -ceu - | consul_ip=$(cat /workspace/consul-ip) - image="${_REGION}-docker.pkg.dev/$PROJECT_ID/${_REPOSITORY}/select-ai:${_IMAGE_TAG}" - mtls_args=() + image="${_IMAGE_URI}" + if [[ -z "$image" ]]; then + image="${_REGION}-docker.pkg.dev/$PROJECT_ID/${_REPOSITORY}/select-ai:${_IMAGE_TAG}" + fi + secret_mappings=("SELECT_AI_DB_CONNECT_STRING=${_DB_DSN_SECRET}:latest") 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" + secret_mappings+=( + "/var/run/select-ai-mtls-ca/ca.pem=select-ai-a2a-mtls-ca:latest" + "/var/run/select-ai-mtls-cert/tls.crt=select-ai-a2a-server-mtls-cert:latest" + "/var/run/select-ai-mtls-key/tls.key=select-ai-a2a-server-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}" \ + secret_mappings_csv=$(IFS=,; echo "${secret_mappings[*]}") + oauth_arg="" + if [[ "${_REQUIRE_OAUTH}" == "true" ]]; then + oauth_arg=",--require-oauth" + fi + gcloud run deploy "${_SERVER_SERVICE}" \ --image="$image" \ --region="${_REGION}" \ + --service-account="${_SERVER_RUNTIME_SA}" \ --no-allow-unauthenticated \ --network="${_NETWORK}" \ --subnet="${_SUBNET}" \ --vpc-egress=all-traffic \ - --max-instances=1 \ + --max-instances=20 \ --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=${_SESSION_TTL_SECONDS}$mtls_env" \ - "${mtls_args[@]}" - gateway_url=$(gcloud run services describe "${_GATEWAY_SERVICE}" \ + --args=a2a,serve,--deployment,clustered,--host,0.0.0.0$oauth_arg \ + --set-env-vars="PUBLIC_URL=https://pending.invalid,SELECT_AI_A2A_TEAM=${_A2A_TEAM},CONSUL_HTTP_URL=http://$consul_ip:8500,WORKER_SERVICE=select-ai-a2a-worker,SESSION_TTL_SECONDS=${_SESSION_TTL_SECONDS}$mtls_env" \ + --set-secrets="$secret_mappings_csv" + server_url=$(gcloud run services describe "${_SERVER_SERVICE}" \ --region="${_REGION}" --format='value(status.url)') - gcloud run services update "${_GATEWAY_SERVICE}" \ + gcloud run services update "${_SERVER_SERVICE}" \ --region="${_REGION}" \ - --update-env-vars="AGENT_URL=$gateway_url" - -images: - - ${_REGION}-docker.pkg.dev/$PROJECT_ID/${_REPOSITORY}/select-ai:${_IMAGE_TAG} + --update-env-vars="PUBLIC_URL=$server_url" substitutions: _REGION: us-central1 - _CLUSTER: select-ai-a2a-gateway + _CLUSTER: select-ai-a2a-cluster _REPOSITORY: select-ai _IMAGE_TAG: dev - _GATEWAY_SERVICE: select-ai-a2a-gateway + _IMAGE_URI: "" + _SERVER_SERVICE: select-ai-a2a-cluster + _SERVER_RUNTIME_SA: unused@example.invalid + _REQUIRE_OAUTH: "false" + _A2A_TEAM: ORACLE_AI_DATABASE_AGENT + _DB_DSN_SECRET: select-ai-a2a-cluster-db-connect-string _NETWORK: default _SUBNET: default _WORKER_REPLICAS: "2" _SESSION_TTL_SECONDS: "900" _ENABLE_WORKER_MTLS: "false" _ROTATE_WORKER_MTLS: "false" - _GKE_DNS_DOMAIN: select-ai-a2a-gateway.internal + _GKE_DNS_DOMAIN: select-ai-a2a-cluster.internal options: logging: CLOUD_LOGGING_ONLY diff --git a/gcloud/gateway/deploy.sh b/gcloud/cluster/deploy.sh similarity index 67% rename from gcloud/gateway/deploy.sh rename to gcloud/cluster/deploy.sh index 08a93ab..1e6c8cc 100755 --- a/gcloud/gateway/deploy.sh +++ b/gcloud/cluster/deploy.sh @@ -7,32 +7,38 @@ # https://oss.oracle.com/licenses/upl. # ----------------------------------------------------------------------------- -# Build and deploy the dynamic Select AI gateway stack: Cloud Run gateway plus +# Build and deploy clustered Select AI A2A: one public Cloud Run server plus # Consul and worker replicas in GKE. set -euo pipefail usage() { cat <<'EOF' -Usage: gcloud/gateway/deploy.sh [options] +Usage: gcloud/cluster/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) + --cluster NAME GKE Autopilot cluster name (default: select-ai-a2a-cluster) --gke-dns-domain DOMAIN Unique GKE additive VPC DNS domain - (default: select-ai-a2a-gateway.internal) + (default: select-ai-a2a-cluster.internal) --repository NAME Artifact Registry Docker repository (default: select-ai) - --gateway-service NAME Cloud Run gateway service name (default: select-ai-a2a-gateway) + --service NAME Cloud Run service name (default: select-ai-a2a-cluster) + --require-oauth Require end-user OAuth bearer authentication + --a2a-team TEAM Database AI team fixed by deployment + (default: ORACLE_AI_DATABASE_AGENT) + --db-dsn-secret NAME Secret containing the fixed database DSN + --rotate-db-dsn Prompt for and rotate the database DSN --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 + --enable-worker-mtls Use ephemeral mTLS certificates for server-to-worker calls --rotate-worker-mtls Replace the existing worker mTLS CA and certificates --mtls-cert-validity-days DAYS - Gateway and worker certificate lifetime (default: 365) + Server and worker certificate lifetime (default: 365) + --image-uri URI Reuse this image for the server and workers --image-tag TAG Image tag (default: git SHA plus UTC timestamp) -h, --help Show this help EOF @@ -41,10 +47,14 @@ EOF 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" +cluster="select-ai-a2a-cluster" +gke_dns_domain="select-ai-a2a-cluster.internal" repository="select-ai" -gateway_service="select-ai-a2a-gateway" +server_service="select-ai-a2a-cluster" +require_oauth="false" +a2a_team="ORACLE_AI_DATABASE_AGENT" +db_dsn_secret="" +rotate_db_dsn="false" network="default" subnet="default" worker_replicas="2" @@ -53,6 +63,7 @@ enable_worker_mtls="false" rotate_worker_mtls="false" mtls_cert_validity_days="365" image_tag="" +image_uri="" while [[ $# -gt 0 ]]; do case "$1" in @@ -61,7 +72,11 @@ while [[ $# -gt 0 ]]; do --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 ;; + --service) server_service="${2:?--service requires a value}"; shift 2 ;; + --require-oauth) require_oauth="true"; shift ;; + --a2a-team) a2a_team="${2:?--a2a-team requires a value}"; shift 2 ;; + --db-dsn-secret) db_dsn_secret="${2:?--db-dsn-secret requires a value}"; shift 2 ;; + --rotate-db-dsn) rotate_db_dsn="true"; shift ;; --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 ;; @@ -76,6 +91,7 @@ while [[ $# -gt 0 ]]; do shift 2 ;; --image-tag) image_tag="${2:?--image-tag requires a value}"; shift 2 ;; + --image-uri) image_uri="${2:?--image-uri requires a value}"; shift 2 ;; -h|--help) usage; exit 0 ;; *) echo "Unknown option: $1" >&2; usage >&2; exit 2 ;; esac @@ -88,6 +104,7 @@ 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 +db_dsn_secret="${db_dsn_secret:-${server_service}-db-connect-string}" if ! [[ "$worker_replicas" =~ ^[1-9][0-9]*$ ]]; then echo "--worker-replicas must be a positive integer." >&2 exit 2 @@ -108,8 +125,14 @@ if ! [[ "$mtls_cert_validity_days" =~ ^[1-9][0-9]*$ ]]; then echo "--mtls-cert-validity-days must be a positive integer." >&2 exit 2 fi +if [[ -n "$image_uri" && -n "$image_tag" ]]; then + echo "--image-uri and --image-tag cannot be used together." >&2 + exit 2 +fi -image_tag="${image_tag:-$(git -C "$repo_root" rev-parse --short HEAD)-$(date -u +%Y%m%d%H%M%S)}" +if [[ -z "$image_uri" ]]; then + image_tag="${image_tag:-$(git -C "$repo_root" rev-parse --short HEAD)-$(date -u +%Y%m%d%H%M%S)}" +fi gcloud services enable \ artifactregistry.googleapis.com \ @@ -128,6 +151,30 @@ if ! gcloud artifacts repositories describe "$repository" \ --repository-format=docker --location="$region" --project="$project_id" fi +project_number="$(gcloud projects describe "$project_id" \ + --format='value(projectNumber)')" +runtime_sa="${project_number}-compute@developer.gserviceaccount.com" + +if [[ "$rotate_db_dsn" == "true" ]] \ + || ! gcloud secrets describe "$db_dsn_secret" \ + --project="$project_id" >/dev/null 2>&1; then + read -r -p "ADB connect descriptor: " db_dsn + trap 'unset db_dsn' EXIT + if gcloud secrets describe "$db_dsn_secret" \ + --project="$project_id" >/dev/null 2>&1; then + printf %s "$db_dsn" | gcloud secrets versions add "$db_dsn_secret" \ + --project="$project_id" --data-file=- >/dev/null + else + printf %s "$db_dsn" | gcloud secrets create "$db_dsn_secret" \ + --project="$project_id" --replication-policy=automatic \ + --data-file=- >/dev/null + fi +fi +gcloud secrets add-iam-policy-binding "$db_dsn_secret" \ + --project="$project_id" \ + --member="serviceAccount:$runtime_sa" \ + --role="roles/secretmanager.secretAccessor" >/dev/null + if ! gcloud container clusters describe "$cluster" \ --location="$region" --project="$project_id" >/dev/null 2>&1; then gcloud container clusters create-auto "$cluster" \ @@ -144,22 +191,19 @@ else 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" + worker_certificate_dns_name="*.select-ai-a2a-worker.select-ai-a2a.svc.$gke_dns_domain" 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 + select-ai-a2a-mtls-ca \ + select-ai-a2a-server-mtls-cert \ + select-ai-a2a-server-mtls-key \ + select-ai-a2a-worker-mtls-cert \ + select-ai-a2a-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" @@ -170,7 +214,7 @@ if [[ "$enable_worker_mtls" == "true" ]]; then # Read the portable text representation and perform a literal SAN check. existing_worker_certificate_text="$( gcloud secrets versions access latest \ - --secret=select-ai-worker-mtls-cert --project="$project_id" 2>/dev/null | \ + --secret=select-ai-a2a-worker-mtls-cert --project="$project_id" 2>/dev/null | \ openssl x509 -noout -text 2>/dev/null || true )" if ! printf '%s\n' "$existing_worker_certificate_text" | \ @@ -190,37 +234,37 @@ 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=select-ai-worker" >/dev/null 2>&1 + -subj "/CN=select-ai-a2a-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 \ -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" \ + -keyout "$mtls_dir/server.key" -out "$mtls_dir/server.csr" \ + -subj "/CN=select-ai-a2a-server" >/dev/null 2>&1 + printf 'extendedKeyUsage=clientAuth\n' > "$mtls_dir/server.ext" + openssl x509 -req -days "$mtls_cert_validity_days" -in "$mtls_dir/server.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 + -out "$mtls_dir/server.crt" -extfile "$mtls_dir/server.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 + select-ai-a2a-mtls-ca \ + select-ai-a2a-server-mtls-cert \ + select-ai-a2a-server-mtls-key \ + select-ai-a2a-worker-mtls-cert \ + select-ai-a2a-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 \ + gcloud secrets versions add select-ai-a2a-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 \ + gcloud secrets versions add select-ai-a2a-server-mtls-cert \ + --data-file="$mtls_dir/server.crt" --project="$project_id" + gcloud secrets versions add select-ai-a2a-server-mtls-key \ + --data-file="$mtls_dir/server.key" --project="$project_id" + gcloud secrets versions add select-ai-a2a-worker-mtls-cert \ --data-file="$mtls_dir/worker.crt" --project="$project_id" - gcloud secrets versions add select-ai-worker-mtls-key \ + gcloud secrets versions add select-ai-a2a-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." @@ -230,13 +274,12 @@ if [[ "$enable_worker_mtls" == "true" ]]; then else echo "Reusing existing worker mTLS material." fi - runtime_sa="${project_number}-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 + select-ai-a2a-mtls-ca \ + select-ai-a2a-server-mtls-cert \ + select-ai-a2a-server-mtls-key \ + select-ai-a2a-worker-mtls-cert \ + select-ai-a2a-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 @@ -247,7 +290,12 @@ build_substitutions=( "_CLUSTER=$cluster" "_REPOSITORY=$repository" "_IMAGE_TAG=$image_tag" - "_GATEWAY_SERVICE=$gateway_service" + "_IMAGE_URI=$image_uri" + "_SERVER_SERVICE=$server_service" + "_SERVER_RUNTIME_SA=$runtime_sa" + "_REQUIRE_OAUTH=$require_oauth" + "_A2A_TEAM=$a2a_team" + "_DB_DSN_SECRET=$db_dsn_secret" "_NETWORK=$network" "_SUBNET=$subnet" "_WORKER_REPLICAS=$worker_replicas" @@ -260,14 +308,14 @@ build_substitutions=( gcloud builds submit "$repo_root" \ --project="$project_id" \ --region="$region" \ - --config="$repo_root/gcloud/gateway/cloudbuild.yaml" \ + --config="$repo_root/gcloud/cluster/cloudbuild.yaml" \ --substitutions="$(IFS=,; printf '%s' "${build_substitutions[*]}")" -gateway_url="$(gcloud run services describe "$gateway_service" \ +server_url="$(gcloud run services describe "$server_service" \ --region="$region" --project="$project_id" --format='value(status.url)')" gemini_service_agent="service-$project_number@gcp-sa-discoveryengine.iam.gserviceaccount.com" -gcloud run services add-iam-policy-binding "$gateway_service" \ +gcloud run services add-iam-policy-binding "$server_service" \ --region="$region" --project="$project_id" \ --member="serviceAccount:$gemini_service_agent" \ --role="roles/run.invoker" >/dev/null @@ -283,8 +331,8 @@ if gcloud iam service-accounts describe "$active_account" \ else deployer_member="user:$active_account" fi -gcloud run services add-iam-policy-binding "$gateway_service" \ +gcloud run services add-iam-policy-binding "$server_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" +printf 'Clustered Agent Card URL:\n%s/.well-known/agent-card.json\n' "$server_url" diff --git a/gcloud/gateway/gke/consul.yaml b/gcloud/cluster/gke/consul.yaml similarity index 94% rename from gcloud/gateway/gke/consul.yaml rename to gcloud/cluster/gke/consul.yaml index 19569c0..796f1bb 100644 --- a/gcloud/gateway/gke/consul.yaml +++ b/gcloud/cluster/gke/consul.yaml @@ -2,7 +2,7 @@ apiVersion: apps/v1 kind: Deployment metadata: name: consul-server - namespace: select-ai-gateway + namespace: select-ai-a2a spec: replicas: 1 selector: @@ -37,7 +37,7 @@ apiVersion: v1 kind: Service metadata: name: consul-server - namespace: select-ai-gateway + namespace: select-ai-a2a annotations: networking.gke.io/load-balancer-type: Internal spec: diff --git a/gcloud/gateway/gke/namespace.yaml b/gcloud/cluster/gke/namespace.yaml similarity index 61% rename from gcloud/gateway/gke/namespace.yaml rename to gcloud/cluster/gke/namespace.yaml index dc1ff53..5209539 100644 --- a/gcloud/gateway/gke/namespace.yaml +++ b/gcloud/cluster/gke/namespace.yaml @@ -1,4 +1,4 @@ apiVersion: v1 kind: Namespace metadata: - name: select-ai-gateway + name: select-ai-a2a diff --git a/gcloud/gateway/gke/worker-mtls.yaml b/gcloud/cluster/gke/worker-mtls.yaml similarity index 69% rename from gcloud/gateway/gke/worker-mtls.yaml rename to gcloud/cluster/gke/worker-mtls.yaml index bdab56e..3c4dc93 100644 --- a/gcloud/gateway/gke/worker-mtls.yaml +++ b/gcloud/cluster/gke/worker-mtls.yaml @@ -1,22 +1,22 @@ apiVersion: apps/v1 kind: StatefulSet metadata: - name: select-ai-worker - namespace: select-ai-gateway + name: select-ai-a2a-worker + namespace: select-ai-a2a spec: - serviceName: select-ai-worker + serviceName: select-ai-a2a-worker replicas: WORKER_REPLICAS selector: matchLabels: - app: select-ai-worker + app: select-ai-a2a-worker template: metadata: labels: - app: select-ai-worker + app: select-ai-a2a-worker spec: terminationGracePeriodSeconds: 30 containers: - - name: worker + - name: select-ai-a2a-worker image: WORKER_IMAGE command: [select-ai] args: @@ -26,6 +26,12 @@ spec: - 0.0.0.0 - --port - "8443" + - --consul-url + - http://consul-server.select-ai-a2a.svc.cluster.local:8500 + - --worker-id + - $(WORKER_ID) + - --worker-endpoint + - https://$(POD_NAME).select-ai-a2a-worker.select-ai-a2a.svc.GKE_DNS_DOMAIN:8443 - --session-ttl-seconds - "SESSION_TTL_SECONDS" - --tls-cert-file @@ -33,13 +39,11 @@ spec: - --tls-key-file - /var/run/select-ai-mtls/tls.key - --tls-ca-file - - /var/run/select-ai-mtls/gateway-ca.crt + - /var/run/select-ai-mtls/server-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: @@ -52,10 +56,6 @@ spec: valueFrom: fieldRef: fieldPath: status.podIP - - name: WORKER_PORT - value: "8443" - - name: WORKER_ENDPOINT - 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 @@ -65,8 +65,8 @@ spec: 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 + - name: server-client-ca + mountPath: /var/run/select-ai-mtls/server-ca.crt subPath: ca.crt readOnly: true readinessProbe: @@ -78,7 +78,7 @@ spec: volumes: - name: worker-server-tls secret: - secretName: select-ai-worker-server-tls - - name: gateway-client-ca + secretName: select-ai-a2a-worker-server-tls + - name: server-client-ca secret: - secretName: select-ai-gateway-client-ca + secretName: select-ai-a2a-server-client-ca diff --git a/gcloud/gateway/gke/worker-service.yaml b/gcloud/cluster/gke/worker-service.yaml similarity index 80% rename from gcloud/gateway/gke/worker-service.yaml rename to gcloud/cluster/gke/worker-service.yaml index 103fe2a..2a56722 100644 --- a/gcloud/gateway/gke/worker-service.yaml +++ b/gcloud/cluster/gke/worker-service.yaml @@ -1,14 +1,14 @@ apiVersion: v1 kind: Service metadata: - name: select-ai-worker - namespace: select-ai-gateway + name: select-ai-a2a-worker + namespace: select-ai-a2a 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 + app: select-ai-a2a-worker ports: - name: http port: 8080 diff --git a/gcloud/gateway/gke/worker.yaml b/gcloud/cluster/gke/worker.yaml similarity index 72% rename from gcloud/gateway/gke/worker.yaml rename to gcloud/cluster/gke/worker.yaml index 3b5392b..8f22b97 100644 --- a/gcloud/gateway/gke/worker.yaml +++ b/gcloud/cluster/gke/worker.yaml @@ -1,21 +1,21 @@ apiVersion: apps/v1 kind: Deployment metadata: - name: select-ai-worker - namespace: select-ai-gateway + name: select-ai-a2a-worker + namespace: select-ai-a2a spec: replicas: WORKER_REPLICAS selector: matchLabels: - app: select-ai-worker + app: select-ai-a2a-worker template: metadata: labels: - app: select-ai-worker + app: select-ai-a2a-worker spec: terminationGracePeriodSeconds: 30 containers: - - name: worker + - name: select-ai-a2a-worker image: WORKER_IMAGE command: - select-ai @@ -26,14 +26,18 @@ spec: - 0.0.0.0 - --port - "8080" + - --consul-url + - http://consul-server.select-ai-a2a.svc.cluster.local:8500 + - --worker-id + - $(WORKER_ID) + - --worker-endpoint + - http://$(WORKER_ADDRESS):8080 - --session-ttl-seconds - "SESSION_TTL_SECONDS" 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: @@ -42,8 +46,6 @@ spec: valueFrom: fieldRef: fieldPath: status.podIP - - name: WORKER_PORT - value: "8080" readinessProbe: httpGet: path: /health diff --git a/gcloud/standalone/README.md b/gcloud/standalone/README.md index 5c2df66..ecb685d 100644 --- a/gcloud/standalone/README.md +++ b/gcloud/standalone/README.md @@ -65,27 +65,157 @@ gcloud services enable \ --project PROJECT_ID ``` -## Deploy (and update) the A2A server +## Deploy dynamic standalone ```bash -gcloud/standalone/deploy.sh --build +gcloud/standalone/deploy.sh \ + --connection-mode dynamic \ + --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. +With no database options, the deployment fixes neither the Connection URL nor +the AI Agent. The generated A2UI form asks each A2A conversation for Connection +URL, username, password, and AI Agent. The resulting database session runs in +an isolated child process and remains bound to the conversation until it +expires. With `--require-oauth`, that binding also includes the authenticated +owner. Dynamic standalone is restricted to one Cloud Run instance because its +session routing is in memory. + +Use `--db-dsn-secret NAME` to fix the Connection URL from an existing Secret +Manager secret, and `--a2a-team TEAM` to fix the AI Agent. The generated form +contains only the values that are not fixed. Use `--rotate-db-config` to prompt +for a Connection URL and create or rotate the service's default DSN secret. + +| Dynamic deployment options | Generated A2UI fields | +| --- | --- | +| Neither option | Connection URL, Database username, Database password, AI Agent | +| `--db-dsn-secret` only | Database username, Database password, AI Agent | +| `--a2a-team` only | Connection URL, Database username, Database password | +| Both options | Database username, Database password | + +### Fix the Connection URL and AI Agent + +The following command updates the default dynamic service so every A2A +conversation supplies only its database username and password. If the named +secret does not exist, the script prompts for the Connection URL and creates +it. Later invocations reuse the secret without prompting. + +```bash +gcloud/standalone/deploy.sh \ + --project PROJECT_ID \ + --connection-mode dynamic \ + --a2a-team ORACLE_AI_DATABASE_AGENT \ + --db-dsn-secret select-ai-a2a-standalone-dynamic-db-connect-string +``` + +Add `--rotate-db-config` only when the secret already exists and its Connection +URL must be replaced. + +## Choose whether end-user OAuth is required + +By default, the service does not require an application OAuth token. Cloud Run +IAM still allows only authorized callers such as Gemini Enterprise to invoke +the private service. Each A2A conversation gets a separate database session, +but the server does not know or verify which human user owns that conversation. + +To require verified end-user ownership, add `--require-oauth`: + +```bash +gcloud/standalone/deploy.sh \ + --connection-mode dynamic \ + --require-oauth \ + --build +``` + +In this mode every A2A request must include `Authorization: Bearer ...` and +the session is scoped by both the authenticated owner and the conversation. +When registering the agent, configure Gemini Enterprise end-user OAuth rather +than selecting **Skip & Finish**. Gemini Enterprise automatically sends its +Cloud Run invocation identity in `X-Serverless-Authorization`; end-user OAuth +adds the separate `Authorization` header consumed by `a2a serve`. + +The server derives a stable owner from an ID token's `iss` and `sub` claims. +It also accepts opaque OAuth access tokens without storing or logging them, +but a refreshed opaque token starts a new owner/session scope. See +[Register and manage A2A agents](https://docs.cloud.google.com/gemini/enterprise/docs/register-and-manage-an-a2a-agent#configure-authentication-and-authorization) +for the OAuth registration steps. + +### Deploy a separate OAuth-enabled service from an existing image + +A new Cloud Run service cannot inherit an image from another service. Resolve +the most recently updated Select AI image directly from Artifact Registry and +construct an immutable digest URI: + +```bash +PROJECT_ID=PROJECT_ID +REGION=us-central1 +REPOSITORY=select-ai +IMAGE_NAME=select-ai + +IMAGE_RECORD="$(gcloud artifacts docker images list \ + "$REGION-docker.pkg.dev/$PROJECT_ID/$REPOSITORY/$IMAGE_NAME" \ + --project "$PROJECT_ID" \ + --include-tags \ + --sort-by='~UPDATE_TIME' \ + --limit=1 \ + --format='csv[no-heading](package,version)')" + +if [[ -z "$IMAGE_RECORD" ]]; then + echo "No Select AI image exists in Artifact Registry." >&2 + exit 1 +fi + +IMAGE_URI="${IMAGE_RECORD/,/@}" +echo "$IMAGE_URI" +``` + +Then create a separate service that fixes the Connection URL and AI Agent and +requires end-user OAuth: + +```bash +gcloud/standalone/deploy.sh \ + --project "$PROJECT_ID" \ + --region "$REGION" \ + --repository "$REPOSITORY" \ + --connection-mode dynamic \ + --service select-ai-a2a-standalone-oauth \ + --a2a-team ORACLE_AI_DATABASE_AGENT \ + --db-dsn-secret select-ai-a2a-standalone-oauth-db-connect-string \ + --require-oauth \ + --image-uri "$IMAGE_URI" +``` + +On its first invocation, the deployment script prompts for the Connection URL +and creates the service-specific secret. The generated A2UI form contains only +database username and password. Register this service in Gemini Enterprise +with end-user OAuth; selecting **Skip & Finish** causes its requests to fail +with HTTP 401. + +## Deploy fixed standalone + +```bash +gcloud/standalone/deploy.sh \ + --connection-mode fixed \ + --a2a-team ORACLE_AI_DATABASE_AGENT \ + --build +``` + +Fixed mode prompts for the ADB connect descriptor, username, and password. It +mounts all three from Secret Manager and creates the existing shared database +pool at server startup. The generated A2UI connection form is skipped. ### 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): +Wallet configuration is supported only in fixed mode. 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 +gcloud/standalone/deploy.sh \ + --connection-mode fixed \ + --a2a-team ORACLE_AI_DATABASE_AGENT \ + --wallet-archive /path/to/Wallet_database.zip ``` The script prompts for the wallet password, stores the ZIP and password as @@ -104,21 +234,26 @@ deploys private Cloud Run, sets the final public URL in the Agent Card, grants y 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 +The default Cloud Run services are `select-ai-a2a-standalone-dynamic` and +`select-ai-a2a-standalone-fixed`, selected by `--connection-mode`. Dynamic mode +has no default AI Agent. Fixed mode requires `--a2a-team` so all four +connection values are complete at startup. Fix either optional dynamic value with explicit options: ```bash -gcloud/standalone/deploy.sh --service sales-analyst-a2a --a2a-team SALES_ANALYST +gcloud/standalone/deploy.sh \ + --connection-mode dynamic \ + --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. +Use a distinct `--service` value for each fixed deployment profile. 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`. +`--max-instances` controls the number of Cloud Run containers in fixed mode. +Dynamic mode requires exactly one instance. A fixed-mode container can use up +to 10 Oracle connections by default; change that limit with `--pool-max-size`. ### Update the Select AI SDK or this repository @@ -127,20 +262,24 @@ and deploy the new image: ```bash git pull -gcloud/standalone/deploy.sh --build +gcloud/standalone/deploy.sh --connection-mode dynamic --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: +the existing image is reused. Existing explicitly selected database secrets +are reused without prompting. To configure or rotate the dynamic deployment's +Connection URL, explicitly request it: ```bash -gcloud/standalone/deploy.sh --rotate-db-credentials +gcloud/standalone/deploy.sh \ + --connection-mode dynamic \ + --rotate-db-config ``` ### What `cloudbuild.yaml` does -`gcloud/standalone/deploy.sh --build` uses `gcloud/standalone/cloudbuild.yaml` to tell Cloud Build to build +`gcloud/standalone/deploy.sh --connection-mode MODE --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/`. diff --git a/gcloud/standalone/cloudbuild.yaml b/gcloud/standalone/cloudbuild.yaml index 3cedbda..213f0fc 100644 --- a/gcloud/standalone/cloudbuild.yaml +++ b/gcloud/standalone/cloudbuild.yaml @@ -5,7 +5,7 @@ # http://oss.oracle.com/licenses/upl. # ----------------------------------------------------------------------------- -# Build one reusable Select AI image for standalone servers, the gateway, and +# Build one reusable Select AI image for standalone and clustered servers and # GKE worker replicas. Runtime role selection is deployment configuration. steps: - name: gcr.io/cloud-builders/docker diff --git a/gcloud/standalone/deploy.sh b/gcloud/standalone/deploy.sh index 3ac7ff1..e964cfc 100755 --- a/gcloud/standalone/deploy.sh +++ b/gcloud/standalone/deploy.sh @@ -23,13 +23,15 @@ Options: --project PROJECT Google Cloud project (defaults to gcloud config project) --region REGION Cloud Run and Artifact Registry region (default: us-central1) --repository REPOSITORY Docker repository name (default: select-ai) - --service SERVICE Cloud Run service name (default: oracle-a2a-agent) - --a2a-team TEAM Agent Team installed in Oracle Database (default: ORACLE_AI_DATABASE_AGENT) + --connection-mode MODE Required: fixed or dynamic + --require-oauth Require end-user OAuth bearer authentication + --service SERVICE Cloud Run service name (defaults from connection mode) + --a2a-team TEAM Fix the AI Agent in the deployment --runtime-sa EMAIL Runtime service-account email --runtime-sa-name NAME Default runtime service-account name (default: oracle-a2a-runtime) - --db-user-secret NAME Secret name for the ADB user - --db-password-secret NAME Secret name for the ADB password - --db-dsn-secret NAME Secret name for the ADB connect descriptor + --db-user-secret NAME Secret name for the fixed-mode ADB user + --db-password-secret NAME Secret name for the fixed-mode ADB password + --db-dsn-secret NAME Fix the Connection URL using this secret --wallet-secret NAME Secret name for the wallet archive --wallet-password-secret NAME Secret name for the wallet password --wallet-archive PATH Wallet ZIP to upload or replace @@ -40,7 +42,7 @@ Options: --image-uri URI Deploy this container image --build Build the current checkout before deploying --image-tag TAG Tag for --build (default: git SHA plus UTC timestamp) - --rotate-db-credentials Prompt for and rotate ADB credentials + --rotate-db-config Prompt for and rotate the configured database values -h, --help Show this help EOF } @@ -49,14 +51,17 @@ repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" project_id="" region="us-central1" repository="select-ai" -service="oracle-a2a-agent" -a2a_team="ORACLE_AI_DATABASE_AGENT" +connection_mode="" +require_oauth=false +service="" +a2a_team="" runtime_sa="" runtime_sa_name="oracle-a2a-runtime" runtime_sa_explicit=false db_user_secret="" db_password_secret="" db_dsn_secret="" +db_dsn_secret_explicit=false wallet_secret="" wallet_password_secret="" wallet_archive="" @@ -67,20 +72,22 @@ pool_max_size="10" image_uri="" image_tag="" build_image=false -rotate_db_credentials=false +rotate_db_config=false 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 ;; --repository) repository="${2:?--repository requires a value}"; shift 2 ;; + --connection-mode) connection_mode="${2:?--connection-mode requires a value}"; shift 2 ;; + --require-oauth) require_oauth=true; shift ;; --service) service="${2:?--service requires a value}"; shift 2 ;; --a2a-team) a2a_team="${2:?--a2a-team requires a value}"; shift 2 ;; --runtime-sa) runtime_sa="${2:?--runtime-sa requires a value}"; runtime_sa_explicit=true; shift 2 ;; --runtime-sa-name) runtime_sa_name="${2:?--runtime-sa-name requires a value}"; shift 2 ;; --db-user-secret) db_user_secret="${2:?--db-user-secret requires a value}"; shift 2 ;; --db-password-secret) db_password_secret="${2:?--db-password-secret requires a value}"; shift 2 ;; - --db-dsn-secret) db_dsn_secret="${2:?--db-dsn-secret requires a value}"; shift 2 ;; + --db-dsn-secret) db_dsn_secret="${2:?--db-dsn-secret requires a value}"; db_dsn_secret_explicit=true; shift 2 ;; --wallet-secret) wallet_secret="${2:?--wallet-secret requires a value}"; shift 2 ;; --wallet-password-secret) wallet_password_secret="${2:?--wallet-password-secret requires a value}"; shift 2 ;; --wallet-archive) wallet_archive="${2:?--wallet-archive requires a value}"; shift 2 ;; @@ -91,7 +98,7 @@ while [[ $# -gt 0 ]]; do --image-uri) image_uri="${2:?--image-uri requires a value}"; shift 2 ;; --build) build_image=true; shift ;; --image-tag) image_tag="${2:?--image-tag requires a value}"; shift 2 ;; - --rotate-db-credentials) rotate_db_credentials=true; shift ;; + --rotate-db-config) rotate_db_config=true; shift ;; -h|--help) usage; exit 0 ;; *) echo "Unknown option: $1" >&2; usage >&2; exit 2 ;; esac @@ -106,10 +113,56 @@ if [[ -z "$project_id" || "$project_id" == "(unset)" ]]; then exit 1 fi +case "$connection_mode" in + fixed) + service="${service:-select-ai-a2a-standalone-fixed}" + if [[ -z "$a2a_team" ]]; then + echo "Fixed standalone requires --a2a-team." >&2 + exit 2 + fi + ;; + dynamic) + service="${service:-select-ai-a2a-standalone-dynamic}" + if [[ "$max_instances" != "1" ]]; then + echo "Dynamic standalone requires --max-instances 1 because session routing is in memory." >&2 + exit 2 + fi + if [[ -n "$wallet_archive" ]]; then + echo "Dynamic standalone does not currently support --wallet-archive." >&2 + exit 2 + fi + ;; + *) + echo "--connection-mode must be fixed or dynamic." >&2 + exit 2 + ;; +esac + +form_summary="" +if [[ "$connection_mode" == "dynamic" ]]; then + append_form_field() { + if [[ -n "$form_summary" ]]; then + form_summary+=", " + fi + form_summary+="$1" + } + if [[ -z "$db_dsn_secret" && "$rotate_db_config" == false ]]; then + append_form_field "Connection URL" + fi + append_form_field "Database username" + append_form_field "Database password" + if [[ -z "$a2a_team" ]]; then + append_form_field "AI Agent" + fi +fi +form_summary="${form_summary:-none}" + runtime_sa="${runtime_sa:-${runtime_sa_name}@${project_id}.iam.gserviceaccount.com}" db_user_secret="${db_user_secret:-${service}-db-user}" db_password_secret="${db_password_secret:-${service}-db-password}" -db_dsn_secret="${db_dsn_secret:-${service}-db-connect-string}" +if [[ "$connection_mode" == "fixed" || "$db_dsn_secret_explicit" == true || "$rotate_db_config" == true ]]; then + db_dsn_secret="${db_dsn_secret:-${service}-db-connect-string}" +fi wallet_secret="${wallet_secret:-${service}-wallet}" wallet_password_secret="${wallet_password_secret:-${service}-wallet-password}" @@ -150,23 +203,30 @@ if gcloud run services describe "$service" --project="$project_id" --region="$re fi create_or_rotate_secrets=false -if [[ "$rotate_db_credentials" == true ]]; then +if [[ "$rotate_db_config" == true ]]; then create_or_rotate_secrets=true -else +elif [[ "$connection_mode" == "fixed" ]]; then for secret in "$db_user_secret" "$db_password_secret" "$db_dsn_secret"; do if ! gcloud secrets describe "$secret" --project="$project_id" >/dev/null 2>&1; then create_or_rotate_secrets=true break fi done +elif [[ -n "$db_dsn_secret" ]] \ + && ! gcloud secrets describe "$db_dsn_secret" --project="$project_id" >/dev/null 2>&1; then + create_or_rotate_secrets=true fi if [[ "$create_or_rotate_secrets" == true ]]; then - echo "Creating or rotating ADB credentials for Cloud Run service: $service" - read -r -p "ADB user: " db_user - read -r -s -p "ADB password: " db_password - echo + echo "Creating or rotating database configuration for Cloud Run service: $service" read -r -p "ADB connect descriptor: " db_dsn + db_user="" + db_password="" + if [[ "$connection_mode" == "fixed" ]]; then + read -r -p "ADB user: " db_user + read -r -s -p "ADB password: " db_password + echo + fi trap 'unset db_user db_password db_dsn' EXIT add_secret() { @@ -177,20 +237,33 @@ if [[ "$create_or_rotate_secrets" == true ]]; then else printf %s "$value" | gcloud secrets create "$name" --project="$project_id" --replication-policy=automatic --data-file=- >/dev/null fi - gcloud secrets add-iam-policy-binding "$name" --project="$project_id" \ - --member="serviceAccount:$runtime_sa" --role="roles/secretmanager.secretAccessor" >/dev/null } - add_secret "$db_user_secret" "$db_user" - add_secret "$db_password_secret" "$db_password" add_secret "$db_dsn_secret" "$db_dsn" + if [[ "$connection_mode" == "fixed" ]]; then + add_secret "$db_user_secret" "$db_user" + add_secret "$db_password_secret" "$db_password" + fi +fi + +if [[ "$connection_mode" == "fixed" ]]; then + for secret in "$db_user_secret" "$db_password_secret" "$db_dsn_secret"; do + gcloud secrets add-iam-policy-binding "$secret" --project="$project_id" \ + --member="serviceAccount:$runtime_sa" \ + --role="roles/secretmanager.secretAccessor" >/dev/null + done +elif [[ -n "$db_dsn_secret" ]]; then + gcloud secrets add-iam-policy-binding "$db_dsn_secret" --project="$project_id" \ + --member="serviceAccount:$runtime_sa" \ + --role="roles/secretmanager.secretAccessor" >/dev/null fi # An Oracle mTLS wallet is a ZIP archive containing several files, so it is # mounted as a Secret Manager volume rather than exposed as an environment # variable. Pass --wallet-archive to enable or replace this optional configuration. wallet_enabled=false -if gcloud secrets describe "$wallet_secret" --project="$project_id" >/dev/null 2>&1 \ +if [[ "$connection_mode" == "fixed" ]] \ + && gcloud secrets describe "$wallet_secret" --project="$project_id" >/dev/null 2>&1 \ && gcloud secrets describe "$wallet_password_secret" --project="$project_id" >/dev/null 2>&1; then wallet_enabled=true fi @@ -239,25 +312,44 @@ fi # Cloud Run needs a URL before the server can construct its Agent Card. Deploy # once with a placeholder, then update PUBLIC_URL with the assigned URL. -secret_mappings=( - "SELECT_AI_USER=$db_user_secret:latest" - "SELECT_AI_PASSWORD=$db_password_secret:latest" - "SELECT_AI_DB_CONNECT_STRING=$db_dsn_secret:latest" -) +secret_mappings_csv="" +append_secret_mapping() { + if [[ -n "$secret_mappings_csv" ]]; then + secret_mappings_csv+="," + fi + secret_mappings_csv+="$1" +} +if [[ -n "$db_dsn_secret" ]]; then + append_secret_mapping "SELECT_AI_DB_CONNECT_STRING=$db_dsn_secret:latest" +fi +if [[ "$connection_mode" == "fixed" ]]; then + append_secret_mapping "SELECT_AI_USER=$db_user_secret:latest" + append_secret_mapping "SELECT_AI_PASSWORD=$db_password_secret:latest" +fi if [[ "$wallet_enabled" == true ]]; then - secret_mappings+=( - "/var/run/secrets/select-ai-wallet/wallet.zip=$wallet_secret:latest" - "SELECT_AI_WALLET_PASSWORD=$wallet_password_secret:latest" - ) + append_secret_mapping "/var/run/secrets/select-ai-wallet/wallet.zip=$wallet_secret:latest" + append_secret_mapping "SELECT_AI_WALLET_PASSWORD=$wallet_password_secret:latest" +fi +secret_option="--clear-secrets" +if [[ -n "$secret_mappings_csv" ]]; then + secret_option="--set-secrets=$secret_mappings_csv" +fi + +env_vars_csv="PUBLIC_URL=https://pending.invalid,SELECT_AI_POOL_MAX_SIZE=$pool_max_size" +if [[ -n "$a2a_team" ]]; then + env_vars_csv+=",SELECT_AI_A2A_TEAM=$a2a_team" +fi +if [[ "$require_oauth" == true ]]; then + env_vars_csv+=",SELECT_AI_A2A_REQUIRE_OAUTH=true" fi -secret_mappings_csv="$(IFS=,; echo "${secret_mappings[*]}")" gcloud run deploy "$service" --image="$image_uri" --project="$project_id" --region="$region" \ --service-account="$runtime_sa" --no-allow-unauthenticated --port=8080 \ --command="/app/docker/a2a-entrypoint.sh" \ - --memory="$memory" --timeout="$timeout" --max-instances="$max_instances" \ - --set-env-vars="SELECT_AI_A2A_TEAM=$a2a_team,PUBLIC_URL=https://pending.invalid,SELECT_AI_POOL_MAX_SIZE=$pool_max_size" \ - --update-secrets="$secret_mappings_csv" + --memory="$memory" --timeout="$timeout" --min-instances=1 \ + --max-instances="$max_instances" \ + --set-env-vars="$env_vars_csv" \ + "$secret_option" service_url="$(gcloud run services describe "$service" --project="$project_id" --region="$region" --format='value(status.url)')" gcloud run services update "$service" --project="$project_id" --region="$region" --update-env-vars="PUBLIC_URL=$service_url" @@ -301,7 +393,26 @@ Deployment complete. Cloud Run URL: $service_url Gemini Enterprise invoker: $gemini_sa +Connection mode: $connection_mode +End-user OAuth required: $require_oauth +A2UI connection fields: $form_summary + +Agent Card endpoint: + $service_url/.well-known/agent-card.json -Paste this A2A v0.3 Agent Card into Gemini Enterprise: - select-ai a2a agent-card --team "$a2a_team" --public-url "$service_url" EOF +if [[ "$require_oauth" == true ]]; then + cat <<'EOF' + +Gemini Enterprise must be configured with end-user OAuth so it sends an +Authorization bearer token in addition to its automatic +X-Serverless-Authorization Cloud Run identity token. +EOF +else + cat <<'EOF' + +End-user OAuth is not required. Cloud Run IAM authenticates Gemini Enterprise, +and database sessions are separated by A2A conversation rather than by a +verified human-user identity. +EOF +fi diff --git a/samples/README.md b/samples/README.md index c083e41..932b4f9 100644 --- a/samples/README.md +++ b/samples/README.md @@ -219,22 +219,22 @@ Task 7e1...: completed } ``` -## A2A dynamic gateway +## Dynamic A2A sessions -The dynamic gateway samples submit the A2UI database connection form, open a -temporary worker session, and execute database tasks. The gateway advertises +The dynamic session samples submit the A2UI database connection form, open a +temporary database session, and execute database tasks. The server advertises `streaming: false` and supports non-blocking task execution with `configuration.blocking: false` and `tasks/get`. -Gateway-specific samples that perform the form handshake and then execute a -real database task are in [a2a/gateway](a2a/gateway/README.md): +Dynamic-session samples that perform the form handshake and then execute a +real database task are in [a2a/dynamic](a2a/dynamic/README.md): ```bash -python samples/a2a/gateway/blocking_task.py -python samples/a2a/gateway/task_poll.py +python samples/a2a/dynamic/blocking_task.py +python samples/a2a/dynamic/task_poll.py ``` -See that README for local Consul, worker, and gateway startup instructions. +See that README for standalone and clustered startup instructions. The full A2A architecture, protocol details, session lifecycle, and Google Cloud deployment explanation are in the diff --git a/samples/a2a/dynamic/README.md b/samples/a2a/dynamic/README.md new file mode 100644 index 0000000..7a52e1f --- /dev/null +++ b/samples/a2a/dynamic/README.md @@ -0,0 +1,162 @@ +# Dynamic A2A session samples + +These samples connect to either standalone or clustered dynamic deployment, +inspect and submit its A2UI database connection form, and execute Select AI +tasks against the database. + +The dynamic server advertises `streaming: false` and supports request/response task +operations. Long-running work is returned as a task and can be followed with +`tasks/get`. + +The server supports asynchronous work with `message/send` and +`configuration.blocking: false`, followed by `tasks/get`. + +These samples use the A2A v0.3 JSON-RPC names used by the existing samples: +`message/send`, `tasks/get`, and `tasks/cancel`. A client using A2A 1.0 should +send `A2A-Version: 1.0` and use `SendMessage`, `GetTask`, `ListTasks`, and +`CancelTask`; its non-blocking option is `configuration.returnImmediately`. +The server accepts both versions, but streaming is disabled in both. + +The sample scripts import `call`, `connect`, `send_prompt`, and +`print_task_summary` from the adjacent +[`_common.py`](https://github.com/oracle/python-select-ai/blob/main/samples/a2a/dynamic/_common.py) +file. This is a +repository-local sample helper, not an additional Python dependency. It sends +the JSON-RPC requests, reads the fields requested by the returned A2UI form, +submits only those fields from the environment variables below, and formats +the final task result. If you copy a script elsewhere, copy +[`_common.py`](https://github.com/oracle/python-select-ai/blob/main/samples/a2a/dynamic/_common.py) +with it or replace those helpers with your own A2A client code. + +## Local setup + +Install the A2A extra if necessary: + +```bash +source .venv/bin/activate +pip install -e '.[a2a]' +``` + +### Dynamic standalone + +This is the smallest complete local setup. Fix the DSN and team on the server +so the generated form contains only username and password. Ensure database +username and password variables are absent from the server process so it does +not select the fixed shared-pool path: + +```bash +env -u SELECT_AI_USER -u SELECT_AI_PASSWORD \ + select-ai a2a serve \ + --deployment standalone \ + --host 127.0.0.1 \ + --port 8000 \ + --public-url http://127.0.0.1:8000 \ + --dsn '' \ + --team ORACLE_AI_DATABASE_AGENT +``` + +Dynamic standalone keeps its routing in memory and therefore runs as one +server process and one service instance. It does not require Consul or the +separate worker command. + +### Clustered + +Run these commands in three terminals from the repository root. + +Terminal 1, Consul: + +```bash +consul agent -dev -bind=127.0.0.1 -client=127.0.0.1 +``` + +Terminal 2, one worker: + +```bash +source .venv/bin/activate + +select-ai a2a worker \ + --host 127.0.0.1 \ + --port 8081 \ + --consul-url http://127.0.0.1:8500 \ + --worker-id local-worker \ + --worker-endpoint http://127.0.0.1:8081 +``` + +Terminal 3, the dynamic server: + +```bash +source .venv/bin/activate + +select-ai a2a serve \ + --deployment clustered \ + --host 127.0.0.1 \ + --port 8000 \ + --public-url http://127.0.0.1:8000 \ + --consul-url http://127.0.0.1:8500 +``` + +For either deployment, export the form values in the terminal that runs the +sample scripts. The helper submits only the properties requested by the +server, so deployment-fixed DSN or team values are not sent back: + +```bash +export SELECT_AI_DB_CONNECT_STRING='' +export SELECT_AI_USER='' +export SELECT_AI_PASSWORD='' +export SELECT_AI_A2A_TEAM='ORACLE_AI_DATABASE_AGENT' +``` + +For a TNS-alias DSN, set `TNS_ADMIN` in the worker terminal before starting +the clustered worker, or in the standalone server terminal. Wallet-based +Oracle Database mTLS is not currently supported by the dynamic session path. + +To require OAuth, start the server with `--require-oauth`. Then set its bearer +token without changing the scripts: + +```bash +export SELECT_AI_A2A_BEARER_TOKEN='' +``` + +Override the default endpoint with `SELECT_AI_A2A_ENDPOINT`. + +## Run the samples + +First inspect the exact A2UI artifact and requested fields without submitting +credentials: + +```bash +python samples/a2a/dynamic/inspect_form.py +``` + +The task samples perform the connection-form handshake automatically and +validate that the final artifact is +`database-agent-result`. + +Blocking database task: + +```bash +python samples/a2a/dynamic/blocking_task.py +``` + +Non-blocking database task with polling: + +```bash +python samples/a2a/dynamic/task_poll.py +``` + +Expected task output is similar to: + +```text +Task : completed +Artifact: database-agent-result +... +``` + +For a quick health check: + +```bash +curl http://127.0.0.1:8081/health +curl -sS http://127.0.0.1:8000/.well-known/agent-card.json | jq +``` + +The worker health endpoint applies only to clustered deployment. diff --git a/samples/a2a/gateway/_common.py b/samples/a2a/dynamic/_common.py similarity index 59% rename from samples/a2a/gateway/_common.py rename to samples/a2a/dynamic/_common.py index 25b03a5..2d3d9ab 100644 --- a/samples/a2a/gateway/_common.py +++ b/samples/a2a/dynamic/_common.py @@ -5,7 +5,7 @@ # https://oss.oracle.com/licenses/upl. # ----------------------------------------------------------------------------- -"""Small A2A v0.3 client helpers for the dynamic gateway samples.""" +"""Small A2A v0.3 client helpers for dynamic session samples.""" import json import os @@ -13,14 +13,23 @@ from urllib.request import Request, urlopen ENDPOINT = os.environ.get( - "SELECT_AI_A2A_GATEWAY_ENDPOINT", + "SELECT_AI_A2A_ENDPOINT", "http://127.0.0.1:8000/a2a/jsonrpc/", ) TEAM_NAME = os.environ.get("SELECT_AI_A2A_TEAM", "ORACLE_AI_DATABASE_AGENT") +_CONNECTION_ENV = { + "connection_url": "SELECT_AI_DB_CONNECT_STRING", + "username": "SELECT_AI_USER", + "password": "SELECT_AI_PASSWORD", + "ai_agent": "SELECT_AI_A2A_TEAM", +} def call(method: str, params: dict) -> dict: """Make one A2A v0.3 JSON-RPC call and return its result.""" + headers = {"Content-Type": "application/json"} + if token := os.environ.get("SELECT_AI_A2A_BEARER_TOKEN"): + headers["Authorization"] = f"Bearer {token}" request = Request( ENDPOINT, data=json.dumps( @@ -31,7 +40,7 @@ def call(method: str, params: dict) -> dict: "params": params, } ).encode(), - headers={"Content-Type": "application/json"}, + headers=headers, method="POST", ) with urlopen(request) as response: # noqa: S310 @@ -53,13 +62,20 @@ def _message(text: str, context_id: str | None = None) -> dict: def connect(prompt: str) -> str: - """Bootstrap one gateway session and return its context ID.""" - form_task = call("message/send", {"message": _message(prompt)}) - if form_task["artifacts"][0]["name"] != "database-connection-form": - raise RuntimeError( - "Expected database-connection-form, got " - f"{form_task['artifacts'][0].get('name')}" - ) + """Submit the requested form fields and return the connected context.""" + form_task, fields = request_connection_form(prompt) + submitted = {} + for field in fields: + env_name = _CONNECTION_ENV[field] + value = os.environ.get(env_name) + if field == "ai_agent" and not value: + value = TEAM_NAME + if not value: + raise RuntimeError( + f"The connection form requires {field}; set {env_name}." + ) + + submitted[field] = value context_id = form_task["contextId"] connection_task = call( @@ -77,16 +93,7 @@ def connect(prompt: str) -> str: "version": "v0.9", "action": { "name": "submit_database_connection", - "context": { - "dsn": os.environ[ - "SELECT_AI_DB_CONNECT_STRING" - ], - "username": os.environ["SELECT_AI_USER"], - "password": os.environ[ - "SELECT_AI_PASSWORD" - ], - "team_name": TEAM_NAME, - }, + "context": submitted, }, }, "metadata": {"mimeType": "application/json+a2ui"}, @@ -101,12 +108,47 @@ def connect(prompt: str) -> str: return context_id +def request_connection_form(prompt: str) -> tuple[dict, tuple[str, ...]]: + """Request and inspect the server-generated A2UI connection form.""" + form_task = call("message/send", {"message": _message(prompt)}) + if form_task["artifacts"][0]["name"] != "database-connection-form": + raise RuntimeError( + "Expected database-connection-form, got " + f"{form_task['artifacts'][0].get('name')}" + ) + + fields = _connection_form_fields(form_task) + if not fields: + raise RuntimeError("The A2UI connection form contains no fields.") + return form_task, fields + + +def _connection_form_fields(form_task: dict) -> tuple[str, ...]: + """Read canonical fields from the A2UI submit action.""" + artifact = form_task["artifacts"][0] + for part in artifact.get("parts") or []: + operation = part.get("data") or {} + update = operation.get("updateComponents") or {} + for component in update.get("components") or []: + event = (component.get("action") or {}).get("event") or {} + if event.get("name") == "submit_database_connection": + context = event.get("context") or {} + unsupported = set(context) - set(_CONNECTION_ENV) + if unsupported: + raise RuntimeError( + "Connection form requested unsupported fields: " + + ", ".join(sorted(unsupported)) + ) + return tuple(context) + return () + + def send_prompt( prompt: str, context_id: str, blocking: bool | None = None, ) -> dict: - """Send a normal database prompt through an existing gateway session.""" + """Send a normal database prompt through an existing dynamic session.""" params = {"message": _message(prompt, context_id)} if blocking is False: params["configuration"] = {"blocking": False} diff --git a/samples/a2a/gateway/blocking_task.py b/samples/a2a/dynamic/blocking_task.py similarity index 89% rename from samples/a2a/gateway/blocking_task.py rename to samples/a2a/dynamic/blocking_task.py index 5f18f8a..9d92307 100644 --- a/samples/a2a/gateway/blocking_task.py +++ b/samples/a2a/dynamic/blocking_task.py @@ -5,7 +5,7 @@ # https://oss.oracle.com/licenses/upl. # ----------------------------------------------------------------------------- -"""Connect to a dynamic A2A gateway, then send a blocking database request.""" +"""Connect to a dynamic A2A server, then send a blocking database request.""" from _common import connect, print_task_summary, send_prompt diff --git a/samples/a2a/dynamic/inspect_form.py b/samples/a2a/dynamic/inspect_form.py new file mode 100644 index 0000000..f4622fa --- /dev/null +++ b/samples/a2a/dynamic/inspect_form.py @@ -0,0 +1,17 @@ +# ----------------------------------------------------------------------------- +# 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. +# ----------------------------------------------------------------------------- + +"""Request and print the dynamic A2UI database connection form.""" + +import json + +from _common import request_connection_form + +task, fields = request_connection_form("Show the database connection form.") +print(f"Context: {task['contextId']}") +print("Requested fields: " + ", ".join(fields)) +print(json.dumps(task["artifacts"][0], indent=2)) diff --git a/samples/a2a/gateway/task_poll.py b/samples/a2a/dynamic/task_poll.py similarity index 94% rename from samples/a2a/gateway/task_poll.py rename to samples/a2a/dynamic/task_poll.py index 698611c..b675a09 100644 --- a/samples/a2a/gateway/task_poll.py +++ b/samples/a2a/dynamic/task_poll.py @@ -5,7 +5,7 @@ # https://oss.oracle.com/licenses/upl. # ----------------------------------------------------------------------------- -"""Connect to a dynamic A2A gateway, then poll a database task.""" +"""Connect to a dynamic A2A server, then poll a database task.""" import time diff --git a/samples/a2a/gateway/README.md b/samples/a2a/gateway/README.md deleted file mode 100644 index d5056f4..0000000 --- a/samples/a2a/gateway/README.md +++ /dev/null @@ -1,118 +0,0 @@ -# Dynamic A2A gateway samples - -These samples connect to a dynamic gateway, submit its A2UI database -connection form, open a temporary worker session, and execute Select AI tasks -against the database. - -The gateway advertises `streaming: false` and supports request/response task -operations. Long-running work is returned as a task and can be followed with -`tasks/get`. - -The gateway supports asynchronous work with `message/send` and -`configuration.blocking: false`, followed by `tasks/get`. - -These samples use the A2A v0.3 JSON-RPC names used by the existing samples: -`message/send`, `tasks/get`, and `tasks/cancel`. A client using A2A 1.0 should -send `A2A-Version: 1.0` and use `SendMessage`, `GetTask`, `ListTasks`, and -`CancelTask`; its non-blocking option is `configuration.returnImmediately`. -The gateway accepts both versions, but streaming is disabled in both. - -The sample scripts import `call`, `connect`, `send_prompt`, and -`print_task_summary` from the adjacent -[`_common.py`](https://github.com/oracle/python-select-ai/blob/main/samples/a2a/gateway/_common.py) -file. This is a -repository-local sample helper, not an additional Python dependency. It sends -the JSON-RPC requests, performs the A2UI connection-form handshake using the -environment variables below, and formats the final task result. If you copy a -script elsewhere, copy -[`_common.py`](https://github.com/oracle/python-select-ai/blob/main/samples/a2a/gateway/_common.py) -with it or replace those helpers with your own A2A client code. - -## Local setup - -Install the A2A extra if necessary: - -```bash -source .venv/bin/activate -pip install -e '.[a2a]' -``` - -Run these commands in three terminals from the repository root. - -Terminal 1, Consul: - -```bash -consul agent -dev -bind=127.0.0.1 -client=127.0.0.1 -``` - -Terminal 2, one worker: - -```bash -source .venv/bin/activate - -CONSUL_HTTP_URL=http://127.0.0.1:8500 \ -WORKER_ID=local-worker \ -WORKER_ADDRESS=127.0.0.1 \ -WORKER_PORT=8081 \ -select-ai a2a worker --host 127.0.0.1 --port 8081 -``` - -Terminal 3, the gateway: - -```bash -source .venv/bin/activate - -select-ai a2a gateway \ - --host 127.0.0.1 \ - --port 8000 \ - --agent-url http://127.0.0.1:8000 \ - --consul-url http://127.0.0.1:8500 -``` - -The worker must be able to connect to the database when a sample submits the -form. Export the same values used by the other samples, plus the optional -team name: - -```bash -export SELECT_AI_DB_CONNECT_STRING='' -export SELECT_AI_USER='' -export SELECT_AI_PASSWORD='' -export SELECT_AI_A2A_TEAM='ORACLE_AI_DATABASE_AGENT' -``` - -For a TNS-alias DSN, set `TNS_ADMIN` in the worker terminal before starting -the worker. Wallet-based Oracle Database mTLS is not currently supported by -the gateway session connection path. - -## Run the samples - -The gateway-specific samples perform the connection-form handshake -automatically and validate that the final artifact is -`database-agent-result`. - -Blocking database task: - -```bash -python samples/a2a/gateway/blocking_task.py -``` - -Non-blocking database task with polling: - -```bash -python samples/a2a/gateway/task_poll.py -``` - -Expected task output is similar to: - -```text -Task : completed -Artifact: database-agent-result -... -``` - -For a quick health check: - -```bash -curl http://127.0.0.1:8081/health -curl -sS http://127.0.0.1:8000/.well-known/agent-card.json | jq -``` diff --git a/src/select_ai/agent/a2a/__init__.py b/src/select_ai/agent/a2a/__init__.py index 7424546..0517008 100644 --- a/src/select_ai/agent/a2a/__init__.py +++ b/src/select_ai/agent/a2a/__init__.py @@ -7,7 +7,14 @@ """A2A support for Select AI Agent Teams and temporary sessions.""" -from .models import GatewaySettings, SessionInfo, SessionRoute +from .models import ( + ConnectionConfig, + GatewaySettings, + SessionInfo, + SessionRoute, + StandaloneSessionSettings, + WorkerSettings, +) def create_gateway_app(settings): @@ -17,20 +24,30 @@ def create_gateway_app(settings): return factory(settings) -def create_worker_app( - session_ttl_seconds: int = 900, - session_start_timeout_seconds: int = 30, -): +def create_worker_app(settings: WorkerSettings): """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) + return factory(settings) + + +def create_embedded_session_app(settings: StandaloneSessionSettings): + """Build a dynamic standalone A2A/A2UI application.""" + from select_ai.agent.a2a.embedded import ( + create_embedded_session_app as factory, + ) + + return factory(settings) __all__ = [ "GatewaySettings", + "ConnectionConfig", "SessionInfo", "SessionRoute", + "StandaloneSessionSettings", + "WorkerSettings", "create_gateway_app", + "create_embedded_session_app", "create_worker_app", ] diff --git a/src/select_ai/agent/a2a/auth.py b/src/select_ai/agent/a2a/auth.py new file mode 100644 index 0000000..44265f7 --- /dev/null +++ b/src/select_ai/agent/a2a/auth.py @@ -0,0 +1,147 @@ +# ----------------------------------------------------------------------------- +# 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. +# ----------------------------------------------------------------------------- + +"""Authentication integration for public A2A Starlette applications.""" + +from __future__ import annotations + +import base64 +import binascii +import hashlib +import json + +from a2a.types import ( + AgentCard, + HTTPAuthSecurityScheme, + SecurityScheme, + StringList, +) +from starlette.authentication import ( + AuthCredentials, + AuthenticationBackend, + AuthenticationError, + SimpleUser, +) +from starlette.middleware import Middleware +from starlette.middleware.authentication import AuthenticationMiddleware +from starlette.responses import JSONResponse + + +class TrustedBearerAuthenticationBackend(AuthenticationBackend): + """Map an upstream-validated OAuth bearer token to a Starlette user. + + Signature and claims validation belongs to the deployment's authenticating + proxy (for example Gemini Enterprise). This backend only establishes the + SDK user object from the already trusted token. + """ + + async def authenticate(self, connection): + header = connection.headers.get("authorization") + if header is None: + return None + scheme, separator, token = header.partition(" ") + if separator != " " or scheme.lower() != "bearer" or not token: + raise AuthenticationError("Authorization must use Bearer syntax.") + owner = _bearer_owner(token) + return AuthCredentials(["authenticated"]), SimpleUser(owner) + + +class RequireA2AAuthenticationMiddleware: + """Require an SDK-visible user on the JSON-RPC endpoint.""" + + def __init__(self, app) -> None: + self.app = app + + async def __call__(self, scope, receive, send) -> None: + if ( + scope["type"] == "http" + and scope.get("path", "").rstrip("/") == "/a2a/jsonrpc" + and not scope["user"].is_authenticated + ): + response = JSONResponse( + {"detail": "Bearer authentication is required."}, + status_code=401, + headers={"WWW-Authenticate": "Bearer"}, + ) + await response(scope, receive, send) + return + await self.app(scope, receive, send) + + +def authentication_middleware(require_oauth: bool) -> list[Middleware]: + """Enable bearer authentication only when explicitly requested.""" + if not require_oauth: + return [] + return [ + Middleware( + AuthenticationMiddleware, + backend=TrustedBearerAuthenticationBackend(), + on_error=_authentication_error, + ), + Middleware( + RequireA2AAuthenticationMiddleware, + ), + ] + + +def add_bearer_security(agent_card: AgentCard) -> None: + """Advertise the bearer contract on an authenticated Agent Card.""" + agent_card.security_schemes["bearer"].CopyFrom( + SecurityScheme( + http_auth_security_scheme=HTTPAuthSecurityScheme( + description=( + "Upstream-validated end-user OAuth 2.0 access token or " + "OpenID Connect ID token." + ), + scheme="bearer", + ) + ) + ) + requirement = agent_card.security_requirements.add() + requirement.schemes["bearer"].CopyFrom(StringList()) + + +def _authentication_error(_connection, error): + return JSONResponse( + {"detail": str(error)}, + status_code=401, + headers={"WWW-Authenticate": "Bearer"}, + ) + + +def _bearer_owner(token: str) -> str: + """Derive an opaque owner from a trusted JWT or OAuth access token.""" + jwt_owner = _jwt_owner(token) + if jwt_owner is not None: + return jwt_owner + digest = hashlib.sha256(token.encode("utf-8")).hexdigest() + return f"oauth:{digest}" + + +def _jwt_owner(token: str) -> str | None: + """Return the issuer/subject owner when the bearer value is a JWT.""" + parts = token.split(".") + if len(parts) != 3: + return None + try: + encoded = parts[1] + "=" * (-len(parts[1]) % 4) + payload = json.loads(base64.urlsafe_b64decode(encoded)) + except ( + binascii.Error, + UnicodeDecodeError, + ValueError, + TypeError, + ): + return None + subject = payload.get("sub") if isinstance(payload, dict) else None + issuer = payload.get("iss") if isinstance(payload, dict) else None + if not isinstance(subject, str) or not subject: + return None + if not isinstance(issuer, str) or not issuer: + return None + digest = hashlib.sha256(f"{issuer}\0{subject}".encode("utf-8")).hexdigest() + return f"jwt:{digest}" diff --git a/src/select_ai/agent/a2a/context_store.py b/src/select_ai/agent/a2a/context_store.py index 6dacc37..74ba492 100644 --- a/src/select_ai/agent/a2a/context_store.py +++ b/src/select_ai/agent/a2a/context_store.py @@ -20,15 +20,15 @@ _CREATE_TABLE = """ BEGIN EXECUTE IMMEDIATE ' - CREATE TABLE SELECT_AI_A2A_CONTEXTS ( + CREATE TABLE DBMS_AI_A2A_CONTEXTS$ ( owner VARCHAR2(512) NOT NULL, context_id VARCHAR2(255) NOT NULL, conversation_id VARCHAR2(255) NOT NULL, created_at TIMESTAMP WITH TIME ZONE NOT NULL, - CONSTRAINT select_ai_a2a_contexts_pk PRIMARY KEY (owner, context_id) + CONSTRAINT dbms_ai_a2a_contexts_pk PRIMARY KEY (owner, context_id) )'; EXECUTE IMMEDIATE ' - COMMENT ON TABLE SELECT_AI_A2A_CONTEXTS + COMMENT ON TABLE DBMS_AI_A2A_CONTEXTS$ IS ''Managed by select_ai.a2a.context_store'''; EXCEPTION WHEN OTHERS THEN @@ -38,6 +38,13 @@ END; """ +_CREATE_VIEW = """ + CREATE OR REPLACE VIEW DBMS_AI_A2A_CONTEXTS AS + SELECT owner, context_id, conversation_id, created_at + FROM DBMS_AI_A2A_CONTEXTS$ + WITH READ ONLY +""" + class OracleContextStore: """Persist one Oracle conversation for each A2A context.""" @@ -58,6 +65,7 @@ async def initialize(self) -> None: if self.initialized: return await self._execute(_CREATE_TABLE) + await self._execute(_CREATE_VIEW) self.initialized = True async def get_or_create( @@ -83,7 +91,7 @@ async def get_or_create( try: await self._execute( """ - INSERT INTO SELECT_AI_A2A_CONTEXTS ( + INSERT INTO DBMS_AI_A2A_CONTEXTS$ ( owner, context_id, conversation_id, created_at ) VALUES ( :owner, :context_id, :conversation_id, SYSTIMESTAMP @@ -106,7 +114,7 @@ async def _get(self, owner: str, context_id: str) -> Optional[str]: row = await self._fetchone( """ SELECT conversation_id - FROM SELECT_AI_A2A_CONTEXTS + FROM DBMS_AI_A2A_CONTEXTS$ WHERE owner = :owner AND context_id = :context_id """, owner=owner, diff --git a/src/select_ai/agent/a2a/embedded.py b/src/select_ai/agent/a2a/embedded.py new file mode 100644 index 0000000..105ba7d --- /dev/null +++ b/src/select_ai/agent/a2a/embedded.py @@ -0,0 +1,293 @@ +# ----------------------------------------------------------------------------- +# 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. +# ----------------------------------------------------------------------------- + +"""Embedded session routing for dynamic standalone A2A deployment.""" + +from __future__ import annotations + +import asyncio +from concurrent.futures import Future +from contextlib import asynccontextmanager +from threading import Event, Lock, Thread +from uuid import uuid4 + +from a2a.types.a2a_pb2 import ( + CancelTaskRequest, + GetTaskRequest, + ListTasksRequest, + ListTasksResponse, + Message, + SendMessageRequest, + Task, +) + +from select_ai.agent.a2a.forms import validate_connection_form +from select_ai.agent.a2a.gateway import create_session_app +from select_ai.agent.a2a.models import ( + SessionInfo, + StandaloneSessionSettings, +) +from select_ai.agent.a2a.session_process import ( + ProcessSessionBackend, + SessionNotFound, + SessionSpec, + SessionUnavailable, +) +from select_ai.agent.a2a.worker_client import ReconnectRequired +from select_ai.agent.a2a.worker_protocol import ( + A2AMethod, + decode_result, +) + + +class EmbeddedSessionClient: + """Expose the clustered client contract over an in-process router.""" + + def __init__(self, settings: StandaloneSessionSettings) -> None: + self.backend = ProcessSessionBackend( + settings.session_ttl_seconds, + settings.session_start_timeout_seconds, + ) + self._bindings: dict[tuple[str, str], str] = {} + self._tasks: dict[tuple[str, str], str] = {} + self._routing_lock = Lock() + self._loop: asyncio.AbstractEventLoop | None = None + self._thread: Thread | None = None + self._reaper: Future | None = None + + def start(self) -> None: + """Start the private event loop that owns the async backend.""" + if self._loop is not None: + return + ready = Event() + + def run_loop() -> None: + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + self._loop = loop + ready.set() + loop.run_forever() + loop.close() + + self._thread = Thread( + target=run_loop, + name="select-ai-a2a-sessions", + daemon=True, + ) + self._thread.start() + ready.wait() + self._reaper = asyncio.run_coroutine_threadsafe( + self.backend.reap_expired(), + self._require_loop(), + ) + + def shutdown(self) -> None: + """Close every child process and stop the private event loop.""" + if self._loop is None: + return + if self._reaper is not None: + self._reaper.cancel() + self._run(self.backend.close_all()) + loop = self._require_loop() + loop.call_soon_threadsafe(loop.stop) + if self._thread is not None: + self._thread.join(timeout=10) + self._loop = None + self._thread = None + self._reaper = None + with self._routing_lock: + self._bindings.clear() + self._tasks.clear() + + def open_session( + self, + owner: str, + context_id: str, + session_info: SessionInfo, + ) -> str: + """Create one isolated child and atomically bind it to a context.""" + session_id = str(uuid4()) + self._run( + self.backend.open( + SessionSpec( + session_id=session_id, + owner=owner, + dsn=session_info.dsn, + username=session_info.username, + password=session_info.password, + team_name=session_info.team_name, + ) + ) + ) + key = (owner, context_id) + with self._routing_lock: + created = key not in self._bindings + if created: + self._bindings[key] = session_id + if not created: + self._run(self.backend.close(session_id)) + raise RuntimeError("A database session already exists.") + return session_id + + def session_exists(self, owner: str, context_id: str) -> bool: + """Return whether the owner/context binding still has a live child.""" + session_id = self._session_id(owner, context_id) + if session_id is None: + return False + try: + self._run(self.backend.get(session_id)) + except SessionNotFound: + self._drop_binding(owner, context_id, session_id) + return False + return True + + def send_message( + self, + owner: str, + context_id: str, + request: SendMessageRequest, + ) -> Task | Message | None: + result = self._dispatch( + owner, + context_id, + A2AMethod.SEND_MESSAGE, + request, + ) + value = decode_result(result) + if isinstance(value, Task): + with self._routing_lock: + self._tasks[(owner, value.id)] = context_id + return value + + def get_task(self, owner: str, request: GetTaskRequest) -> Task | None: + context_id = self._task_context(owner, request.id) + if context_id is None: + return None + value = decode_result( + self._dispatch(owner, context_id, A2AMethod.GET_TASK, request) + ) + return value if isinstance(value, Task) else None + + def list_tasks( + self, + owner: str, + context_id: str, + request: ListTasksRequest, + ) -> ListTasksResponse: + return ( + decode_result( + self._dispatch( + owner, context_id, A2AMethod.LIST_TASKS, request + ) + ) + or ListTasksResponse() + ) + + def cancel_task(self, owner: str, task_id: str) -> Task | None: + context_id = self._task_context(owner, task_id) + if context_id is None: + return None + value = decode_result( + self._dispatch( + owner, + context_id, + A2AMethod.CANCEL_TASK, + CancelTaskRequest(id=task_id), + ) + ) + return value if isinstance(value, Task) else None + + def close_session(self, owner: str, context_id: str) -> None: + """Close and forget one owner/context binding.""" + with self._routing_lock: + session_id = self._bindings.pop((owner, context_id), None) + if session_id is None: + return + try: + self._run(self.backend.close(session_id)) + except SessionNotFound: + pass + + def _dispatch(self, owner, context_id, method, request): + session_id = self._session_id(owner, context_id) + if session_id is None: + raise ReconnectRequired( + "Database session expired; reconnect required.", + context_id, + ) + try: + return self._run( + self.backend.dispatch( + session_id, + method.value, + request.SerializeToString(), + ) + ) + except (SessionNotFound, SessionUnavailable) as error: + self._drop_binding(owner, context_id, session_id) + raise ReconnectRequired(str(error), context_id) from error + + def _session_id(self, owner: str, context_id: str) -> str | None: + with self._routing_lock: + return self._bindings.get((owner, context_id)) + + def _task_context(self, owner: str, task_id: str) -> str | None: + with self._routing_lock: + return self._tasks.get((owner, task_id)) + + def _drop_binding( + self, + owner: str, + context_id: str, + session_id: str, + ) -> None: + with self._routing_lock: + key = (owner, context_id) + if self._bindings.get(key) == session_id: + self._bindings.pop(key, None) + + def _run(self, coroutine): + return asyncio.run_coroutine_threadsafe( + coroutine, + self._require_loop(), + ).result() + + def _require_loop(self) -> asyncio.AbstractEventLoop: + if self._loop is None: + raise RuntimeError("Embedded session manager is not running.") + return self._loop + + +def create_embedded_session_app( + settings: StandaloneSessionSettings, +): + """Build a dynamic standalone app backed by local child processes.""" + form_template = settings.connection_form_template + if form_template is not None: + form_template = validate_connection_form( + list(form_template), + settings.connection.missing_fields, + ) + client = EmbeddedSessionClient(settings) + + @asynccontextmanager + async def lifespan(_app): + client.start() + try: + yield + finally: + await asyncio.to_thread(client.shutdown) + + return create_session_app( + public_url=settings.public_url, + session_client=client, + connection=settings.connection, + connection_form_template=form_template, + require_oauth=settings.require_oauth, + description=settings.description, + lifespan=lifespan, + ) diff --git a/src/select_ai/agent/a2a/forms.py b/src/select_ai/agent/a2a/forms.py index 9d82bff..7593d06 100644 --- a/src/select_ai/agent/a2a/forms.py +++ b/src/select_ai/agent/a2a/forms.py @@ -5,20 +5,58 @@ # https://oss.oracle.com/licenses/upl. # ----------------------------------------------------------------------------- -"""A2UI connection form emitted by the public gateway.""" +"""Generation and validation for the A2UI database connection form.""" +from __future__ import annotations + +import json +from collections.abc import Iterable +from copy import deepcopy +from pathlib import Path from uuid import uuid4 from select_ai.agent.a2a.a2ui import A2UI_CATALOG_ID, A2UI_VERSION +from select_ai.agent.a2a.models import CONNECTION_FORM_FIELDS + +_ACTION_NAME = "submit_database_connection" +_FIELD_COMPONENTS = { + "connection_url": ("Connection URL", "shortText"), + "username": ("Database username", "shortText"), + "password": ("Database password", "obscured"), + "ai_agent": ("AI Agent", "shortText"), +} -def connection_form(surface_id: str | None = None) -> list[dict]: - """Return the non-persistent database connection form.""" - # A2UI surface IDs must be globally unique for the renderer's lifetime. - # Gemini retains surfaces for an A2A conversation after the connection - # form is submitted, so reusing a fixed ID prevents a reconnect form from - # being created in that same conversation. +def connection_form( + surface_id: str | None = None, + missing_fields: Iterable[str] | None = None, + template: tuple[dict, ...] | None = None, +) -> list[dict]: + """Return a fresh form containing only missing connection properties.""" + fields = ( + tuple(CONNECTION_FORM_FIELDS) + if missing_fields is None + else tuple(missing_fields) + ) + if not fields: + return [] + if set(fields) - set(CONNECTION_FORM_FIELDS): + raise ValueError("Connection form contains unsupported fields.") surface_id = surface_id or f"db-connect-{uuid4().hex}" + if template is not None: + return _rebind_surface(template, surface_id) + + field_components = [ + { + "id": field, + "component": "TextField", + "label": _FIELD_COMPONENTS[field][0], + "value": {"path": f"/{field}"}, + "variant": _FIELD_COMPONENTS[field][1], + } + for field in fields + ] + action_context = {field: {"path": f"/{field}"} for field in fields} return [ { "version": A2UI_VERSION, @@ -36,14 +74,7 @@ def connection_form(surface_id: str | None = None) -> list[dict]: { "id": "column", "component": "Column", - "children": [ - "title", - "dsn", - "user", - "password", - "team", - "connect", - ], + "children": ["title", *fields, "connect"], }, { "id": "title", @@ -51,34 +82,7 @@ def connection_form(surface_id: str | None = None) -> list[dict]: "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", - }, + *field_components, { "id": "connect_label", "component": "Text", @@ -91,13 +95,8 @@ def connection_form(surface_id: str | None = None) -> list[dict]: "variant": "primary", "action": { "event": { - "name": "submit_database_connection", - "context": { - "dsn": {"path": "/dsn"}, - "username": {"path": "/username"}, - "password": {"path": "/password"}, - "team_name": {"path": "/team_name"}, - }, + "name": _ACTION_NAME, + "context": action_context, } }, }, @@ -109,12 +108,111 @@ def connection_form(surface_id: str | None = None) -> list[dict]: "updateDataModel": { "surfaceId": surface_id, "path": "/", - "value": { - "dsn": "", - "username": "", - "password": "", - "team_name": "", - }, + "value": {field: "" for field in fields}, }, }, ] + + +def load_connection_form( + path: str, + missing_fields: Iterable[str], +) -> tuple[dict, ...]: + """Load and validate one custom A2UI form template.""" + try: + value = json.loads(Path(path).read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as error: + raise ValueError(f"Could not load A2UI form: {error}") from error + return validate_connection_form(value, missing_fields) + + +def validate_connection_form( + operations, + missing_fields: Iterable[str], +) -> tuple[dict, ...]: + """Validate the supported custom connection-form contract.""" + fields = tuple(missing_fields) + if not isinstance(operations, list) or not operations: + raise ValueError("A2UI form must be a non-empty JSON array.") + if not all(isinstance(operation, dict) for operation in operations): + raise ValueError("Every A2UI form operation must be an object.") + if any( + operation.get("version") != A2UI_VERSION for operation in operations + ): + raise ValueError(f"Every A2UI operation must use {A2UI_VERSION}.") + + surface_ids = set() + actions = [] + has_catalog = False + password_is_obscured = False + for operation in operations: + for value in operation.values(): + if isinstance(value, dict) and value.get("surfaceId"): + surface_ids.add(value["surfaceId"]) + created = operation.get("createSurface") + if isinstance(created, dict): + has_catalog = created.get("catalogId") == A2UI_CATALOG_ID + updated = operation.get("updateComponents") + if isinstance(updated, dict): + for component in updated.get("components") or []: + if not isinstance(component, dict): + continue + event = (component.get("action") or {}).get("event") + if ( + isinstance(event, dict) + and event.get("name") == _ACTION_NAME + ): + actions.append(event) + if ( + component.get("component") == "TextField" + and component.get("value") == {"path": "/password"} + and component.get("variant") == "obscured" + ): + password_is_obscured = True + data_model = operation.get("updateDataModel") + if isinstance(data_model, dict): + value = data_model.get("value") + if isinstance(value, dict) and value.get("password"): + raise ValueError("A2UI form must not embed a password value.") + + if not has_catalog: + raise ValueError( + "A2UI form must create the advertised catalog surface." + ) + if len(surface_ids) != 1: + raise ValueError( + "A2UI form must use one consistent template surfaceId." + ) + if len(actions) != 1: + raise ValueError( + f"A2UI form must contain exactly one {_ACTION_NAME} action." + ) + context = actions[0].get("context") + if not isinstance(context, dict) or set(context) != set(fields): + raise ValueError( + "A2UI form action must submit exactly the missing fields." + ) + if any(context[field] != {"path": f"/{field}"} for field in fields): + raise ValueError( + "A2UI form action uses an invalid connection field path." + ) + if "password" in fields and not password_is_obscured: + raise ValueError("A2UI password input must use the obscured variant.") + return tuple(deepcopy(operations)) + + +def _rebind_surface(template: tuple[dict, ...], surface_id: str) -> list[dict]: + operations = deepcopy(list(template)) + + def replace(value) -> None: + if isinstance(value, dict): + if "surfaceId" in value: + value["surfaceId"] = surface_id + for child in value.values(): + replace(child) + elif isinstance(value, list): + for child in value: + replace(child) + + replace(operations) + return operations diff --git a/src/select_ai/agent/a2a/gateway.py b/src/select_ai/agent/a2a/gateway.py index 1fc9834..a511319 100644 --- a/src/select_ai/agent/a2a/gateway.py +++ b/src/select_ai/agent/a2a/gateway.py @@ -19,6 +19,7 @@ new_text_part, ) from a2a.server.context import ServerCallContext +from a2a.server.owner_resolver import resolve_user_scope from a2a.server.request_handlers import RequestHandler from a2a.server.routes import create_jsonrpc_routes from a2a.types import ( @@ -58,8 +59,15 @@ a2ui_part, find_action, ) -from select_ai.agent.a2a.forms import connection_form -from select_ai.agent.a2a.models import GatewaySettings, SessionInfo +from select_ai.agent.a2a.auth import ( + add_bearer_security, + authentication_middleware, +) +from select_ai.agent.a2a.forms import ( + connection_form, + validate_connection_form, +) +from select_ai.agent.a2a.models import ConnectionConfig, GatewaySettings from select_ai.agent.a2a.worker_client import ReconnectRequired, WorkerClient from select_ai.version import __version__ @@ -94,9 +102,19 @@ def _unsupported_stream(*_args, **_kwargs): class GatewayRequestHandler(RequestHandler): """Proxy A2A requests to the Oracle-backed handler in a worker child.""" - def __init__(self, worker_client: WorkerClient, agent_card: AgentCard): + def __init__( + self, + worker_client: WorkerClient, + agent_card: AgentCard, + connection: ConnectionConfig | None = None, + connection_form_template: tuple[dict, ...] | None = None, + require_oauth: bool = False, + ): self.worker_client = worker_client self.agent_card = agent_card + self.connection = connection or ConnectionConfig() + self.connection_form_template = connection_form_template + self.require_oauth = require_oauth # RequestHandler requires every operation even when the Agent Card does # not advertise streaming or push notifications. @@ -110,14 +128,15 @@ def __init__(self, worker_client: WorkerClient, agent_card: AgentCard): async def on_message_send( self, params: SendMessageRequest, - _context: ServerCallContext, + context: ServerCallContext, ) -> Task | Message: """Handle connection bootstrap or forward the request to a worker.""" message = params.message + owner = self._owner(context) if action := find_action(message, _CONNECTION_ACTION_NAME): - return await self._handle_connection_action(message, action) + return await self._handle_connection_action(message, action, owner) - response = await self._recover_task_context(message) + response = await self._recover_task_context(message, owner) if response is not None: return response @@ -126,20 +145,29 @@ async def on_message_send( if not await asyncio.to_thread( self.worker_client.session_exists, + owner, context_id, ): - return self._build_connection_form_task( - task_id=message.task_id or None, - context_id=context_id, - history=None if had_task_id else [message], - ) + if self._connection_config.missing_fields: + return self._build_connection_form_task( + task_id=message.task_id or None, + context_id=context_id, + history=None if had_task_id else [message], + ) + if await self._open_session({}, owner, context_id) is None: + return self._build_connection_failure_task( + task_id=message.task_id or None, + context_id=context_id, + history=None if had_task_id else [message], + ) - return await self._forward_message(context_id, params) + return await self._forward_message(owner, context_id, params) async def _handle_connection_action( self, message: Message, action: dict, + owner: str, ) -> Task: """Open a database session from the submitted A2UI form.""" if not message.context_id: @@ -153,20 +181,25 @@ async def _handle_connection_action( ) session_id = await self._open_session( action.get("context") or {}, + owner, message.context_id, ) - text = ( - _CONNECTION_SUCCESS_MESSAGE - if session_id is not None - else _CONNECTION_FAILURE_MESSAGE - ) + if session_id is None: + return self._build_connection_failure_task( + task_id=task.id, + context_id=task.context_id, + ) return self._complete_task( task, - [new_text_part(text)], + [new_text_part(_CONNECTION_SUCCESS_MESSAGE)], "database-session", ) - async def _recover_task_context(self, message: Message) -> Task | None: + async def _recover_task_context( + self, + message: Message, + owner: str, + ) -> Task | None: """Recover a context when a client sends only a previous task ID.""" if message.context_id or not message.task_id: return None @@ -174,6 +207,7 @@ async def _recover_task_context(self, message: Message) -> Task | None: try: existing_task = await asyncio.to_thread( self.worker_client.get_task, + owner, GetTaskRequest(id=message.task_id), ) except ReconnectRequired as error: @@ -206,6 +240,7 @@ def _ensure_context_id(message: Message) -> str: async def _forward_message( self, + owner: str, context_id: str, params: SendMessageRequest, ) -> Task | Message: @@ -214,6 +249,7 @@ async def _forward_message( try: result = await asyncio.to_thread( self.worker_client.send_message, + owner, context_id, params, ) @@ -229,6 +265,7 @@ async def _forward_message( params.message.ClearField("task_id") result = await asyncio.to_thread( self.worker_client.send_message, + owner, context_id, params, ) @@ -245,12 +282,14 @@ async def _forward_message( async def _open_session( self, action_context: dict, + owner: str, context_id: str, ) -> str | None: try: - session_info = SessionInfo.from_a2ui_event(action_context) + session_info = self._connection_config.resolve(action_context) session_id = await asyncio.to_thread( self.worker_client.open_session, + owner, context_id, session_info, ) @@ -261,11 +300,12 @@ async def _open_session( async def on_get_task( self, params: GetTaskRequest, - _context: ServerCallContext, + context: ServerCallContext, ) -> Task | None: try: task = await asyncio.to_thread( self.worker_client.get_task, + self._owner(context), params, ) except ReconnectRequired as error: @@ -285,7 +325,7 @@ async def on_get_task( async def on_list_tasks( self, params: ListTasksRequest, - _context: ServerCallContext, + context: ServerCallContext, ) -> ListTasksResponse: if not params.context_id: raise InvalidParamsError( @@ -294,6 +334,7 @@ async def on_list_tasks( try: return await asyncio.to_thread( self.worker_client.list_tasks, + self._owner(context), params.context_id, params, ) @@ -303,11 +344,12 @@ async def on_list_tasks( async def on_cancel_task( self, params: CancelTaskRequest, - _context: ServerCallContext, + context: ServerCallContext, ) -> Task | None: try: task = await asyncio.to_thread( self.worker_client.cancel_task, + self._owner(context), params.id, ) except ReconnectRequired as error: @@ -331,15 +373,17 @@ async def on_get_extended_agent_card( ) -> AgentCardMessage: return self.agent_card - @staticmethod - def _session_expired_error(context_id: str) -> InvalidParamsError: + def _session_expired_error(self, context_id: str) -> InvalidParamsError: """Build the reconnect error returned by the task-list operation.""" return InvalidParamsError( "Database session expired. Reconnect using contextId.", data={ "reason": "SESSION_EXPIRED", "contextId": context_id, - "connectionForm": connection_form(), + "connectionForm": connection_form( + missing_fields=self._connection_config.missing_fields, + template=self._form_template, + ), }, ) @@ -375,6 +419,12 @@ def _build_connection_form_task( history: list[Message] | None = None, ) -> Task: """Build a completed, transient task containing the connection form.""" + if not self._connection_config.missing_fields: + return self._build_connection_failure_task( + task_id=task_id, + context_id=context_id, + history=history, + ) task = _new_bootstrap_task( task_id=task_id, context_id=context_id, @@ -382,14 +432,61 @@ def _build_connection_form_task( ) return self._complete_task( task, - [ - a2ui_part(item) - for item in connection_form(f"db-connect-{task.id}") - ], + self._connection_form_parts(task.id), "database-connection-form", [A2UI_EXTENSION_URI], ) + def _build_connection_failure_task( + self, + *, + task_id: str | None, + context_id: str, + history: list[Message] | None = None, + ) -> Task: + """Return a safe failure and a fresh form when input is still needed.""" + task = _new_bootstrap_task( + task_id=task_id, + context_id=context_id, + history=history, + ) + parts = [new_text_part(_CONNECTION_FAILURE_MESSAGE)] + parts.extend(self._connection_form_parts(task.id)) + extensions = [A2UI_EXTENSION_URI] if len(parts) > 1 else None + return self._complete_task( + task, + parts, + "database-connection-failure", + extensions, + ) + + def _connection_form_parts(self, task_id: str) -> list: + return [ + a2ui_part(item) + for item in connection_form( + surface_id=f"db-connect-{task_id}", + missing_fields=self._connection_config.missing_fields, + template=self._form_template, + ) + ] + + @property + def _connection_config(self) -> ConnectionConfig: + return getattr(self, "connection", ConnectionConfig()) + + @property + def _form_template(self) -> tuple[dict, ...] | None: + return getattr(self, "connection_form_template", None) + + def _owner(self, context: ServerCallContext) -> str: + """Resolve an OAuth owner or the conversation-only session scope.""" + owner = resolve_user_scope(context) + if owner: + return owner + if not getattr(self, "require_oauth", False): + return "a2a-conversation" + raise PermissionError("Bearer authentication is required.") + def _new_bootstrap_task( *, @@ -408,10 +505,39 @@ def _new_bootstrap_task( 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/" + form_template = settings.connection_form_template + if form_template is not None: + form_template = validate_connection_form( + list(form_template), + settings.connection.missing_fields, + ) + return create_session_app( + public_url=settings.public_url, + session_client=WorkerClient(settings), + connection=settings.connection, + connection_form_template=form_template, + require_oauth=settings.require_oauth, + description=settings.description, + ) + + +def create_session_app( + *, + public_url: str, + session_client, + connection: ConnectionConfig, + connection_form_template: tuple[dict, ...] | None, + require_oauth: bool, + description: str | None = None, + lifespan=None, +) -> Starlette: + """Build the common dynamic-session A2A application.""" + description = description or ( + "Connects a user to a temporary Select AI database session." + ) + endpoint = f"{public_url.rstrip('/')}/a2a/jsonrpc/" card = AgentCard( - name="Select AI Database Gateway", + name="Select AI Database Agent", description=description, version=__version__, default_input_modes=["text/plain", A2UI_MIME_TYPE], @@ -444,7 +570,15 @@ def create_gateway_app(settings: GatewaySettings) -> Starlette: ) ], ) - handler = GatewayRequestHandler(WorkerClient(settings), card) + if require_oauth: + add_bearer_security(card) + handler = GatewayRequestHandler( + session_client, + card, + connection, + connection_form_template, + require_oauth, + ) compat_card = to_compat_agent_card(card).model_dump( by_alias=True, exclude_none=True, @@ -463,4 +597,8 @@ async def get_agent_card(_request): enable_v0_3_compat=True, ) ) - return Starlette(routes=routes) + return Starlette( + routes=routes, + middleware=authentication_middleware(require_oauth), + lifespan=lifespan, + ) diff --git a/src/select_ai/agent/a2a/models.py b/src/select_ai/agent/a2a/models.py index 2812f61..57eca76 100644 --- a/src/select_ai/agent/a2a/models.py +++ b/src/select_ai/agent/a2a/models.py @@ -7,20 +7,113 @@ """Configuration and transient request models for the A2A runtime.""" -from dataclasses import dataclass +from dataclasses import dataclass, field + +CONNECTION_FIELDS = ("dsn", "username", "password", "team_name") +CONNECTION_FORM_FIELDS = ( + "connection_url", + "username", + "password", + "ai_agent", +) +FORM_FIELD_TO_CONNECTION_FIELD = dict( + zip(CONNECTION_FORM_FIELDS, CONNECTION_FIELDS) +) + + +@dataclass(frozen=True) +class ConnectionConfig: + """Deployment-provided values for a database session.""" + + dsn: str | None = None + username: str | None = None + password: str | None = None + team_name: str | None = None + + def __post_init__(self) -> None: + for name in CONNECTION_FIELDS: + value = getattr(self, name) + if value is not None and (not isinstance(value, str) or not value): + raise ValueError(f"{name} must be a non-empty string") + + @property + def missing_fields(self) -> tuple[str, ...]: + """Return canonical A2UI properties that must come from the form.""" + return tuple( + form_name + for form_name, connection_name in FORM_FIELD_TO_CONNECTION_FIELD.items() + if getattr(self, connection_name) is None + ) + + def resolve(self, submitted: dict) -> "SessionInfo": + """Merge validated submitted values with immutable server values.""" + if not isinstance(submitted, dict): + raise ValueError("Connection form context must be an object.") + unknown = set(submitted) - set(CONNECTION_FORM_FIELDS) + if unknown: + raise ValueError("Connection form contains unsupported fields.") + configured = { + form_name + for form_name, connection_name in FORM_FIELD_TO_CONNECTION_FIELD.items() + if getattr(self, connection_name) is not None + } + if configured.intersection(submitted): + raise ValueError( + "Configured connection fields cannot be overridden." + ) + values = { + connection_name: getattr(self, connection_name, None) + or submitted.get(form_name) + for form_name, connection_name in FORM_FIELD_TO_CONNECTION_FIELD.items() + } + return SessionInfo.from_values(values) + + +@dataclass(frozen=True) +class WorkerSettings: + """Configuration for the internal session worker.""" + + consul_url: str + worker_id: str + worker_address: str + worker_port: int + session_ttl_seconds: int + session_start_timeout_seconds: int + worker_endpoint: str | None = None + + def __post_init__(self) -> None: + if self.worker_port < 1 or self.worker_port > 65_535: + raise ValueError("worker_port must be between 1 and 65535") + if self.session_ttl_seconds < 1: + raise ValueError("session_ttl_seconds must be at least 1") + if self.session_start_timeout_seconds < 1: + raise ValueError( + "session_start_timeout_seconds must be at least 1" + ) + object.__setattr__(self, "consul_url", self.consul_url.rstrip("/")) + if self.worker_endpoint: + object.__setattr__( + self, + "worker_endpoint", + self.worker_endpoint.rstrip("/"), + ) @dataclass(frozen=True) class GatewaySettings: """Configuration for the public gateway process.""" - agent_url: str + public_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 + description: str | None = None + connection: ConnectionConfig = field(default_factory=ConnectionConfig) + connection_form_template: tuple[dict, ...] | None = None + require_oauth: bool = False def __post_init__(self) -> None: if self.session_ttl_seconds < 1: @@ -35,7 +128,7 @@ def __post_init__(self) -> None: "worker mTLS requires a CA file, client certificate, and " "client key." ) - object.__setattr__(self, "agent_url", self.agent_url.rstrip("/")) + object.__setattr__(self, "public_url", self.public_url.rstrip("/")) object.__setattr__(self, "consul_url", self.consul_url.rstrip("/")) @property @@ -44,6 +137,28 @@ def worker_mtls_enabled(self) -> bool: return self.worker_tls_ca_file is not None +@dataclass(frozen=True) +class StandaloneSessionSettings: + """Configuration for dynamic sessions embedded in one server process.""" + + public_url: str + session_ttl_seconds: int + session_start_timeout_seconds: int = 30 + description: str | None = None + connection: ConnectionConfig = field(default_factory=ConnectionConfig) + connection_form_template: tuple[dict, ...] | None = None + require_oauth: bool = False + + def __post_init__(self) -> None: + if self.session_ttl_seconds < 1: + raise ValueError("session_ttl_seconds must be at least 1") + if self.session_start_timeout_seconds < 1: + raise ValueError( + "session_start_timeout_seconds must be at least 1" + ) + object.__setattr__(self, "public_url", self.public_url.rstrip("/")) + + @dataclass(frozen=True) class SessionInfo: """Credentials used only while opening one in-memory worker session.""" @@ -54,13 +169,14 @@ class SessionInfo: team_name: str @classmethod - def from_a2ui_event(cls, event: dict) -> "SessionInfo": - required = ("dsn", "username", "password", "team_name") + def from_values(cls, values: dict) -> "SessionInfo": + """Build a complete connection after resolution.""" if not all( - isinstance(event.get(key), str) and event[key] for key in required + isinstance(values.get(key), str) and values[key] + for key in CONNECTION_FIELDS ): raise ValueError("All database connection fields are required.") - return cls(**{key: event[key] for key in required}) + return cls(**{key: values[key] for key in CONNECTION_FIELDS}) @dataclass(frozen=True) @@ -69,3 +185,4 @@ class SessionRoute: endpoint: str expires_at: float + session_id: str diff --git a/src/select_ai/agent/a2a/server.py b/src/select_ai/agent/a2a/server.py index d2c06fa..334c4d3 100644 --- a/src/select_ai/agent/a2a/server.py +++ b/src/select_ai/agent/a2a/server.py @@ -24,6 +24,10 @@ 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.auth import ( + add_bearer_security, + authentication_middleware, +) 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 @@ -83,12 +87,15 @@ def create_app( # noqa: PLR0913 wallet_password: Optional[str] = None, description: Optional[str] = None, pool_max_size: int = 10, + require_oauth: bool = False, ) -> Starlette: """Build an A2A JSON-RPC application for one database AI Agent Team.""" if pool_max_size < 1: raise ValueError("pool_max_size must be at least 1") agent_card = _build_agent_card(team_name, public_url, description) + if require_oauth: + add_bearer_security(agent_card) compat_agent_card = _build_v03_agent_card(agent_card) task_store = OracleTaskStore() context_store = OracleContextStore() @@ -138,7 +145,11 @@ async def get_agent_card(request): enable_v0_3_compat=True, ) ) - return Starlette(routes=routes, lifespan=lifespan) + return Starlette( + routes=routes, + lifespan=lifespan, + middleware=authentication_middleware(require_oauth), + ) def _build_agent_card( diff --git a/src/select_ai/agent/a2a/session_process.py b/src/select_ai/agent/a2a/session_process.py new file mode 100644 index 0000000..2f70b83 --- /dev/null +++ b/src/select_ai/agent/a2a/session_process.py @@ -0,0 +1,452 @@ +# ----------------------------------------------------------------------------- +# 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. +# ----------------------------------------------------------------------------- + +"""Reusable child-process backend for isolated database sessions.""" + +from __future__ import annotations + +import asyncio +import contextlib +import logging +import multiprocessing +import time +from dataclasses import dataclass +from multiprocessing.connection import Connection +from typing import Protocol + +from select_ai.agent.a2a.session_runtime import SessionRuntime +from select_ai.agent.a2a.worker_protocol import ( + PipeMessageType, + ResultKind, + WorkerResult, +) + +LOGGER = logging.getLogger(__name__) +_SESSION_REAPER_INTERVAL_SECONDS = 1 + + +@dataclass(frozen=True) +class SessionSpec: + """Complete inputs needed to start one isolated session.""" + + session_id: str + owner: str + dsn: str + username: str + password: str + team_name: str + + +@dataclass +class ChildSession: + """In-memory ownership record for one database-session process.""" + + process: multiprocessing.Process + connection: Connection + expires_at: float + lock: asyncio.Lock + + +class SessionBackend(Protocol): + """Internal contract shared by worker and embedded session backends.""" + + async def open(self, spec: SessionSpec) -> None: ... + + async def dispatch( + self, + session_id: str, + method: str, + payload: bytes, + ) -> WorkerResult: ... + + async def close(self, session_id: str) -> None: ... + + async def close_all(self) -> None: ... + + async def reap_expired(self) -> None: ... + + +class SessionBackendError(RuntimeError): + """Base error raised by an isolated session backend.""" + + +class SessionNotFound(SessionBackendError): + """The requested session does not exist or is no longer live.""" + + +class SessionUnavailable(SessionBackendError): + """The session process failed while handling a command.""" + + +class SessionStartTimeout(SessionBackendError): + """The session process did not become ready before its deadline.""" + + +class SessionLoginError(SessionBackendError): + """The database connection or team initialization failed.""" + + +class SessionCommandError(SessionBackendError): + """The database runtime rejected one A2A command.""" + + +class ProcessSessionBackend: + """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] = {} + + async def open(self, spec: SessionSpec) -> None: + """Start a child runtime and wait until its database pool is ready.""" + parent_connection, child_connection = multiprocessing.Pipe() + credentials = { + "user": spec.username, + "password": spec.password, + "dsn": spec.dsn, + } + process = multiprocessing.Process( + target=_session_process_main, + args=( + child_connection, + credentials, + spec.session_id, + spec.owner, + spec.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, spec) + except Exception: + await self._terminate(session) + raise + previous = self.sessions.pop(spec.session_id, None) + self.sessions[spec.session_id] = session + if previous: + await self._terminate(previous) + + async def get(self, session_id: str) -> ChildSession: + """Return a live session, closing it if its expiry has elapsed.""" + expired_session = None + 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) + expired_session = session + session = None + if expired_session: + await self._terminate(expired_session) + if session is None: + raise SessionNotFound( + "Database session expired; reconnect required." + ) + return session + + async def dispatch( + self, + session_id: str, + method: str, + payload: bytes, + ) -> WorkerResult: + """Run one A2A operation in the process owning this session.""" + session = await self.get(session_id) + response = None + failure = None + async with session.lock: + if session.process.is_alive(): + try: + await asyncio.to_thread( + session.connection.send, + { + "type": PipeMessageType.A2A.value, + "method": method, + "payload": payload, + }, + ) + response = await self._receive( + session, + timeout_seconds=120, + ) + except (EOFError, OSError, TimeoutError) as error: + failure = error + if failure is not None: + await self._discard(session_id, session) + raise SessionUnavailable( + "Database session is unavailable; reconnect required." + ) from failure + if response is None: + await self._discard(session_id, session) + raise SessionNotFound( + "Database session expired; reconnect required." + ) + if response.get("type") == PipeMessageType.RESULT.value: + return response.get( + "result", + WorkerResult(ResultKind.NONE), + ) + if response.get("type") == PipeMessageType.A2A_ERROR.value: + raise SessionCommandError( + response.get("detail", "A2A request failed.") + ) + LOGGER.error("Select AI session process reported a command failure.") + raise SessionUnavailable( + "Database session is unavailable; reconnect required." + ) + + async def close(self, session_id: str) -> None: + """Terminate a session explicitly.""" + session = self.sessions.pop(session_id, None) + if session is None: + raise SessionNotFound( + "Database session expired; reconnect required." + ) + await self._terminate(session) + + async def reap_expired(self) -> None: + """Continuously terminate expired or dead child sessions.""" + while True: + await asyncio.sleep(_SESSION_REAPER_INTERVAL_SECONDS) + await self._reap_expired_sessions() + + async def close_all(self) -> None: + """Terminate all child sessions during backend shutdown.""" + sessions = list(self.sessions.values()) + self.sessions.clear() + for session in sessions: + await self._terminate(session) + + async def _reap_expired_sessions(self) -> None: + now = time.monotonic() + expired = [ + (session_id, session) + for session_id, session in self.sessions.items() + if session.expires_at <= now or not session.process.is_alive() + ] + for session_id, _session in expired: + self.sessions.pop(session_id, None) + + for session_id, session in expired: + try: + await self._terminate(session) + except Exception: + LOGGER.exception( + "Failed to terminate expired session %s.", + session_id, + ) + + async def _wait_ready( + self, + session: ChildSession, + spec: SessionSpec, + ) -> None: + try: + response = await self._receive( + session, + timeout_seconds=self.session_start_timeout_seconds, + ) + except (EOFError, OSError, TimeoutError) as error: + raise SessionStartTimeout( + "Database session start timed out." + ) from error + if response.get("type") == PipeMessageType.READY.value: + return + detail = response.get("detail", "Database login failed.") + for secret in (spec.password, spec.username, spec.dsn): + detail = detail.replace(secret, "[REDACTED]") + LOGGER.error("Select AI session startup failed: %s", detail[-2_000:]) + raise SessionLoginError("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: + 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: + async with session.lock: + try: + with contextlib.suppress(OSError): + await asyncio.to_thread( + session.connection.send, + {"type": PipeMessageType.CLOSE.value}, + ) + await asyncio.to_thread(session.process.join, timeout=5) + for stop in ( + session.process.terminate, + session.process.kill, + ): + if not session.process.is_alive(): + break + await asyncio.to_thread(stop) + await asyncio.to_thread(session.process.join, timeout=5) + finally: + await asyncio.to_thread(session.connection.close) + + +def _session_process_main( + connection: Connection, + credentials: dict[str, str], + session_id: str, + owner: 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, + owner, + team_name, + ) + ) + finally: + connection.close() + + +async def _run_session_process( + connection: Connection, + credentials: dict[str, str], + session_id: str, + owner: str, + team_name: str, +) -> None: + """Open one async connection and execute A2A operations.""" + import select_ai + + runtime: SessionRuntime | None = None + 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.") + runtime = SessionRuntime(session_id, owner, team_name) + await runtime.initialize() + connection.send({"type": PipeMessageType.READY.value}) + ready = True + await _serve_session_commands(connection, runtime) + except Exception as error: + _report_session_process_error(connection, error, not ready) + finally: + await _close_session_process(runtime, select_ai) + + +async def _serve_session_commands( + connection: Connection, + runtime: SessionRuntime, +) -> None: + """Serve commands for one initialized database session.""" + while True: + command = await _receive_session_command(connection) + if ( + command is None + or command.get("type") == PipeMessageType.CLOSE.value + ): + return + if command.get("type") != PipeMessageType.A2A.value: + connection.send( + { + "type": PipeMessageType.ERROR.value, + "detail": "Invalid command.", + } + ) + continue + await _handle_a2a_command(connection, runtime, command) + + +async def _receive_session_command(connection: Connection) -> dict | None: + """Read one command without blocking the event loop.""" + try: + return await asyncio.to_thread(connection.recv) + except EOFError: + return None + + +async def _handle_a2a_command( + connection: Connection, + runtime: SessionRuntime, + command: dict, +) -> None: + """Execute one internal A2A command and send its result.""" + try: + result = await runtime.handle( + command["method"], + command.get("payload", b""), + ) + except Exception as error: + LOGGER.exception("Select AI session A2A command failed") + connection.send( + {"type": PipeMessageType.A2A_ERROR.value, "detail": str(error)} + ) + return + connection.send({"type": PipeMessageType.RESULT.value, "result": result}) + + +def _report_session_process_error( + connection: Connection, + error: Exception, + during_startup: bool, +) -> None: + """Report startup failures without sending errors after readiness.""" + message = ( + "Select AI session process startup failed" + if during_startup + else "Select AI session process failed" + ) + LOGGER.error(message) + if during_startup: + with contextlib.suppress(OSError): + connection.send( + {"type": PipeMessageType.ERROR.value, "detail": str(error)} + ) + + +async def _close_session_process( + runtime: SessionRuntime | None, + select_ai, +) -> None: + """Close the A2A handler and database connection owned by the child.""" + handler = getattr(runtime, "handler", None) + if handler is not None: + with contextlib.suppress(Exception): + await handler.aclose() + with contextlib.suppress(Exception): + await select_ai.async_disconnect() diff --git a/src/select_ai/agent/a2a/session_runtime.py b/src/select_ai/agent/a2a/session_runtime.py index f516557..086b049 100644 --- a/src/select_ai/agent/a2a/session_runtime.py +++ b/src/select_ai/agent/a2a/session_runtime.py @@ -43,8 +43,8 @@ class SessionUser(User): """Internal A2A user used to scope one worker session's database rows.""" - def __init__(self, session_id: str) -> None: - self.session_id = session_id + def __init__(self, owner: str) -> None: + self.owner = owner @property def is_authenticated(self) -> bool: @@ -52,14 +52,15 @@ def is_authenticated(self) -> bool: @property def user_name(self) -> str: - return self.session_id + return self.owner class SessionRuntime: """Own the A2A handler and Oracle stores for one connected database.""" - def __init__(self, session_id: str, team_name: str) -> None: + def __init__(self, session_id: str, owner: str, team_name: str) -> None: self.session_id = session_id + self.owner = owner self.team_name = team_name self.task_store = OracleTaskStore() self.context_store = OracleContextStore() @@ -96,7 +97,7 @@ async def handle( raise RuntimeError("A2A session runtime is not initialized.") operation = A2AMethod(method) - context = ServerCallContext(user=SessionUser(self.session_id)) + context = ServerCallContext(user=SessionUser(self.owner)) if operation == A2AMethod.DELETE_TASK: request = parse_request(GetTaskRequest, payload) await self.task_store.delete(request.id, context) diff --git a/src/select_ai/agent/a2a/task_store.py b/src/select_ai/agent/a2a/task_store.py index 418c145..a98d660 100644 --- a/src/select_ai/agent/a2a/task_store.py +++ b/src/select_ai/agent/a2a/task_store.py @@ -27,16 +27,16 @@ _CREATE_TABLE = """ BEGIN EXECUTE IMMEDIATE ' - CREATE TABLE SELECT_AI_A2A_TASKS ( + CREATE TABLE DBMS_AI_A2A_TASKS$ ( owner VARCHAR2(512) NOT NULL, task_id VARCHAR2(255) NOT NULL, context_id VARCHAR2(255), task_json CLOB NOT NULL CHECK (task_json IS JSON), updated_at TIMESTAMP WITH TIME ZONE NOT NULL, - CONSTRAINT select_ai_a2a_tasks_pk PRIMARY KEY (owner, task_id) + CONSTRAINT dbms_ai_a2a_tasks_pk PRIMARY KEY (owner, task_id) )'; EXECUTE IMMEDIATE ' - COMMENT ON TABLE SELECT_AI_A2A_TASKS + COMMENT ON TABLE DBMS_AI_A2A_TASKS$ IS ''Managed by select_ai.a2a.task_store'''; EXCEPTION WHEN OTHERS THEN @@ -46,6 +46,13 @@ END; """ +_CREATE_VIEW = """ + CREATE OR REPLACE VIEW DBMS_AI_A2A_TASKS AS + SELECT owner, task_id, context_id, task_json, updated_at + FROM DBMS_AI_A2A_TASKS$ + WITH READ ONLY +""" + class OracleTaskStore(TaskStore): """Persist A2A tasks in Oracle Database using Select AI's connection pool.""" @@ -66,6 +73,7 @@ async def initialize(self) -> None: if self.initialized: return await self._execute(_CREATE_TABLE) + await self._execute(_CREATE_VIEW) self.initialized = True async def save(self, task: Task, context: ServerCallContext) -> None: @@ -73,7 +81,7 @@ async def save(self, task: Task, context: ServerCallContext) -> None: await self.initialize() await self._execute( """ - MERGE INTO SELECT_AI_A2A_TASKS target + MERGE INTO DBMS_AI_A2A_TASKS$ target USING ( SELECT :owner AS owner, :task_id AS task_id FROM dual ) source @@ -104,7 +112,7 @@ async def get( row = await self._fetchone( """ SELECT task_json - FROM SELECT_AI_A2A_TASKS + FROM DBMS_AI_A2A_TASKS$ WHERE owner = :owner AND task_id = :task_id """, owner=self._owner(context), @@ -124,7 +132,7 @@ async def list( rows = await self._fetchall( """ SELECT task_json - FROM SELECT_AI_A2A_TASKS + FROM DBMS_AI_A2A_TASKS$ WHERE owner = :owner AND ( :context_id IS NULL OR context_id = :context_id @@ -166,7 +174,7 @@ async def delete(self, task_id: str, context: ServerCallContext) -> None: await self.initialize() await self._execute( """ - DELETE FROM SELECT_AI_A2A_TASKS + DELETE FROM DBMS_AI_A2A_TASKS$ WHERE owner = :owner AND task_id = :task_id """, owner=self._owner(context), diff --git a/src/select_ai/agent/a2a/worker.py b/src/select_ai/agent/a2a/worker.py index 5043955..cb5b972 100644 --- a/src/select_ai/agent/a2a/worker.py +++ b/src/select_ai/agent/a2a/worker.py @@ -5,425 +5,58 @@ # https://oss.oracle.com/licenses/upl. # ----------------------------------------------------------------------------- -"""Internal worker that owns temporary Select AI child processes.""" +"""Internal HTTP and Consul wrapper for isolated session 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, Request, status from pydantic import BaseModel, Field, SecretStr from starlette.responses import Response -from select_ai.agent.a2a.session_runtime import SessionRuntime +from select_ai.agent.a2a.models import WorkerSettings +from select_ai.agent.a2a.session_process import ( + ProcessSessionBackend, + SessionBackendError, + SessionCommandError, + SessionLoginError, + SessionNotFound, + SessionSpec, + SessionStartTimeout, + SessionUnavailable, +) from select_ai.agent.a2a.worker_protocol import ( PROTOBUF_CONTENT_TYPE, WORKER_A2A_METHOD_HEADER, WORKER_RESULT_KIND_HEADER, - PipeMessageType, - ResultKind, - WorkerResult, ) -LOGGER = logging.getLogger(__name__) -_SESSION_REAPER_INTERVAL_SECONDS = 1 - class OpenSessionRequest(BaseModel): """Sensitive request accepted only on the internal worker interface.""" session_id: str + owner: str = Field(min_length=1, max_length=512) 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) - -@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] = {} - - 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 - previous = self.sessions.pop(request.session_id, None) - self.sessions[request.session_id] = session - if previous: - await self._terminate(previous) - - async def get(self, session_id: str) -> ChildSession: - """Return a live session, closing it if its expiry has elapsed.""" - expired_session = None - 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) - expired_session = session - session = None - if expired_session: - await self._terminate(expired_session) - if session is None: - raise HTTPException( - status_code=404, - detail="Database session expired; reconnect required.", - ) - return session - - async def dispatch( - self, - session_id: str, - method: str, - payload: bytes, - ) -> WorkerResult: - """Run one A2A operation in the process owning this session.""" - session = await self.get(session_id) - response = None - failure = None - async with session.lock: - if session.process.is_alive(): - try: - await asyncio.to_thread( - session.connection.send, - { - "type": PipeMessageType.A2A.value, - "method": method, - "payload": payload, - }, - ) - response = await self._receive( - session, - timeout_seconds=120, - ) - except (EOFError, OSError, TimeoutError) as error: - failure = error - if failure is not None: - await self._discard(session_id, session) - raise HTTPException( - status_code=502, - detail=( - "Database session is unavailable; reconnect required." - ), - ) from failure - if response is None: - await self._discard(session_id, session) - raise HTTPException( - status_code=404, - detail="Database session expired; reconnect required.", - ) - if response.get("type") == PipeMessageType.RESULT.value: - return response.get( - "result", - WorkerResult(ResultKind.NONE), - ) - if response.get("type") == PipeMessageType.A2A_ERROR.value: - raise HTTPException( - status_code=400, - detail=response.get("detail", "A2A request failed."), - ) - 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.""" - 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 reap_expired(self) -> None: - """Continuously terminate expired or dead child sessions.""" - while True: - await asyncio.sleep(_SESSION_REAPER_INTERVAL_SECONDS) - await self._reap_expired_sessions() - - async def close_all(self) -> None: - """Terminate all child sessions during worker shutdown.""" - sessions = list(self.sessions.values()) - self.sessions.clear() - for session in sessions: - await self._terminate(session) - - async def _reap_expired_sessions(self) -> None: - now = time.monotonic() - expired = [ - (session_id, session) - for session_id, session in self.sessions.items() - if session.expires_at <= now or not session.process.is_alive() - ] - for session_id, _session in expired: - self.sessions.pop(session_id, None) - - for session_id, session in expired: - try: - await self._terminate(session) - except Exception: - LOGGER.exception( - "Failed to terminate expired session %s.", - session_id, - ) - - 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") == PipeMessageType.READY.value: - 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: - 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: - async with session.lock: - try: - with contextlib.suppress(OSError): - await asyncio.to_thread( - session.connection.send, - {"type": PipeMessageType.CLOSE.value}, - ) - await asyncio.to_thread(session.process.join, timeout=5) - for stop in ( - session.process.terminate, - session.process.kill, - ): - if not session.process.is_alive(): - break - await asyncio.to_thread(stop) - await asyncio.to_thread(session.process.join, timeout=5) - finally: - await asyncio.to_thread(session.connection.close) - - -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 A2A operations.""" - import select_ai - - runtime: SessionRuntime | None = None - 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.") - runtime = SessionRuntime(session_id, team_name) - await runtime.initialize() - connection.send({"type": PipeMessageType.READY.value}) - ready = True - await _serve_session_commands(connection, runtime) - except Exception as error: - _report_session_process_error(connection, error, not ready) - finally: - await _close_session_process(runtime, select_ai) - - -async def _serve_session_commands( - connection: Connection, - runtime: SessionRuntime, -) -> None: - """Serve commands for one initialized database session.""" - while True: - command = await _receive_session_command(connection) - if ( - command is None - or command.get("type") == PipeMessageType.CLOSE.value - ): - return - if command.get("type") != PipeMessageType.A2A.value: - connection.send( - { - "type": PipeMessageType.ERROR.value, - "detail": "Invalid command.", - } - ) - continue - await _handle_a2a_command(connection, runtime, command) - - -async def _receive_session_command(connection: Connection) -> dict | None: - """Read one command without blocking the event loop.""" - try: - return await asyncio.to_thread(connection.recv) - except EOFError: - return None - - -async def _handle_a2a_command( - connection: Connection, - runtime: SessionRuntime, - command: dict, -) -> None: - """Execute one internal A2A command and send its result.""" - try: - result = await runtime.handle( - command["method"], - command.get("payload", b""), + def to_session_spec(self) -> SessionSpec: + """Convert validated HTTP input to the backend-neutral contract.""" + return SessionSpec( + session_id=self.session_id, + owner=self.owner, + dsn=self.dsn, + username=self.username, + password=self.password.get_secret_value(), + team_name=self.team_name, ) - except Exception as error: - LOGGER.exception("Select AI session A2A command failed") - connection.send( - {"type": PipeMessageType.A2A_ERROR.value, "detail": str(error)} - ) - return - connection.send({"type": PipeMessageType.RESULT.value, "result": result}) - - -def _report_session_process_error( - connection: Connection, - error: Exception, - during_startup: bool, -) -> None: - """Report startup failures without sending errors after readiness.""" - message = ( - "Select AI session process startup failed" - if during_startup - else "Select AI session process failed" - ) - LOGGER.error(message) - if during_startup: - with contextlib.suppress(OSError): - connection.send( - {"type": PipeMessageType.ERROR.value, "detail": str(error)} - ) - - -async def _close_session_process( - runtime: SessionRuntime | None, - select_ai, -) -> None: - """Close the A2A handler and database connection owned by the child.""" - handler = getattr(runtime, "handler", None) - if handler is not None: - with contextlib.suppress(Exception): - await handler.aclose() - with contextlib.suppress(Exception): - await select_ai.async_disconnect() async def _register_with_consul( @@ -434,7 +67,7 @@ async def _register_with_consul( worker_endpoint: str | None = None, ) -> None: payload = { - "Name": "select-ai-worker", + "Name": "select-ai-a2a-worker", "ID": worker_id, "Address": worker_address, "Port": worker_port, @@ -477,31 +110,41 @@ async def _deregister_from_consul(consul_url: str, worker_id: str) -> None: ) -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) +def _http_error(error: SessionBackendError) -> HTTPException: + """Map backend-neutral failures to the existing worker HTTP contract.""" + if isinstance(error, SessionNotFound): + status_code = status.HTTP_404_NOT_FOUND + elif isinstance(error, SessionStartTimeout): + status_code = status.HTTP_504_GATEWAY_TIMEOUT + elif isinstance(error, (SessionLoginError, SessionCommandError)): + status_code = status.HTTP_400_BAD_REQUEST + elif isinstance(error, SessionUnavailable): + status_code = status.HTTP_502_BAD_GATEWAY + else: + status_code = status.HTTP_500_INTERNAL_SERVER_ERROR + return HTTPException(status_code=status_code, detail=str(error)) - 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") + +def create_worker_app(settings: WorkerSettings) -> FastAPI: + """Build the internal session-worker HTTP application.""" + backend = ProcessSessionBackend( + settings.session_ttl_seconds, + settings.session_start_timeout_seconds, + ) @asynccontextmanager async def lifespan(_app): await _register_with_consul( - consul_url, - worker_id, - worker_address, - worker_port, - worker_endpoint, + settings.consul_url, + settings.worker_id, + settings.worker_address, + settings.worker_port, + settings.worker_endpoint, + ) + heartbeat = asyncio.create_task( + _heartbeat(settings.consul_url, settings.worker_id) ) - heartbeat = asyncio.create_task(_heartbeat(consul_url, worker_id)) - reaper = asyncio.create_task(worker.reap_expired()) + reaper = asyncio.create_task(backend.reap_expired()) try: yield finally: @@ -511,8 +154,11 @@ async def lifespan(_app): heartbeat.cancel() with contextlib.suppress(asyncio.CancelledError): await heartbeat - await worker.close_all() - await _deregister_from_consul(consul_url, worker_id) + await backend.close_all() + await _deregister_from_consul( + settings.consul_url, + settings.worker_id, + ) app = FastAPI( title="Select AI Session Worker", @@ -527,7 +173,10 @@ async def health() -> dict[str, str]: @app.post("/sessions", status_code=status.HTTP_201_CREATED) async def open_session(request: OpenSessionRequest) -> dict[str, str]: - await worker.open(request) + try: + await backend.open(request.to_session_spec()) + except SessionBackendError as error: + raise _http_error(error) from error return {"status": "opened"} @app.post("/sessions/{session_id}/a2a") @@ -541,11 +190,14 @@ async def handle_a2a( status_code=400, detail=f"Missing {WORKER_A2A_METHOD_HEADER} header.", ) - result = await worker.dispatch( - session_id, - method, - await request.body(), - ) + try: + result = await backend.dispatch( + session_id, + method, + await request.body(), + ) + except SessionBackendError as error: + raise _http_error(error) from error return Response( content=result.payload, media_type=PROTOBUF_CONTENT_TYPE, @@ -559,6 +211,9 @@ async def handle_a2a( status_code=status.HTTP_204_NO_CONTENT, ) async def close_session(session_id: str) -> None: - await worker.close(session_id) + try: + await backend.close(session_id) + except SessionBackendError as error: + raise _http_error(error) from error return app diff --git a/src/select_ai/agent/a2a/worker_client.py b/src/select_ai/agent/a2a/worker_client.py index 3c38b82..71e6df5 100644 --- a/src/select_ai/agent/a2a/worker_client.py +++ b/src/select_ai/agent/a2a/worker_client.py @@ -10,11 +10,13 @@ from __future__ import annotations import base64 +import hashlib import json import logging import time from threading import Lock from urllib.parse import quote +from uuid import uuid4 import requests from a2a.types.a2a_pb2 import ( @@ -40,8 +42,8 @@ LOGGER = logging.getLogger(__name__) -_SESSION_PREFIX = "select-ai/sessions/" -_TASK_PREFIX = "select-ai/tasks/" +_SESSION_PREFIX = "select-ai-a2a/sessions/" +_TASK_PREFIX = "select-ai-a2a/tasks/" class ReconnectRequired(RuntimeError): @@ -69,12 +71,22 @@ def __init__(self, settings: GatewaySettings): ), } - def open_session(self, session_id: str, session_info: SessionInfo) -> str: + def open_session( + self, + owner: str, + context_id: str, + session_info: SessionInfo, + ) -> str: """Open a session and save a non-secret route in Consul.""" + session_id = str(uuid4()) endpoint = self._select_worker() response = requests.post( f"{endpoint}/sessions", - json={"session_id": session_id, **session_info.__dict__}, + json={ + "session_id": session_id, + "owner": owner, + **session_info.__dict__, + }, timeout=45, **self._worker_request_kwargs, ) @@ -82,34 +94,37 @@ def open_session(self, session_id: str, session_info: SessionInfo) -> str: route = SessionRoute( endpoint=endpoint, expires_at=time.time() + self.settings.session_ttl_seconds, + session_id=session_id, ) - if not self._save_route(session_id, route): + if not self._save_route(owner, context_id, route): self._close_worker_session(route, session_id) raise RuntimeError("Could not create the database session.") return session_id def send_message( self, - session_id: str, + owner: str, + context_id: str, request: SendMessageRequest, ) -> Task | Message | None: """Forward one A2A message to the worker session.""" result = self._dispatch( - session_id, + owner, + context_id, A2AMethod.SEND_MESSAGE, request, ) value = decode_result(result) if isinstance(value, Task): - self._save_task_route(value.id, session_id) + self._save_task_route(owner, value.id, context_id) return value - def get_task(self, params: GetTaskRequest) -> Task | None: + def get_task(self, owner: str, params: GetTaskRequest) -> Task | None: """Retrieve a task from its owning worker session.""" - session_id = self.task_session_id(params.id) - if session_id is None: + context_id = self.task_context_id(owner, params.id) + if context_id is None: return None - result = self._dispatch(session_id, A2AMethod.GET_TASK, params) + result = self._dispatch(owner, context_id, A2AMethod.GET_TASK, params) value = decode_result(result) if isinstance(value, Task): return value @@ -117,56 +132,60 @@ def get_task(self, params: GetTaskRequest) -> Task | None: def list_tasks( self, - session_id: str, + owner: str, + context_id: str, params: ListTasksRequest, ) -> ListTasksResponse: """Retrieve one session's database-paginated task response.""" result = self._dispatch( - session_id, + owner, + context_id, A2AMethod.LIST_TASKS, params, ) return decode_result(result) or ListTasksResponse() - def cancel_task(self, task_id: str) -> Task | None: + def cancel_task(self, owner: str, task_id: str) -> Task | None: """Cancel a task in the worker that owns it.""" - session_id = self.task_session_id(task_id) - if session_id is None: + context_id = self.task_context_id(owner, task_id) + if context_id is None: return None result = self._dispatch( - session_id, + owner, + context_id, A2AMethod.CANCEL_TASK, CancelTaskRequest(id=task_id), ) value = decode_result(result) return value if isinstance(value, Task) else None - def delete_task(self, task_id: str) -> None: + def delete_task(self, owner: str, task_id: str) -> None: """Delete a task and its routing metadata.""" - session_id = self.task_session_id(task_id) - if session_id: + context_id = self.task_context_id(owner, task_id) + if context_id: try: self._dispatch( - session_id, + owner, + context_id, A2AMethod.DELETE_TASK, GetTaskRequest(id=task_id), ) except ReconnectRequired: pass - self._delete_task_route(task_id) + self._delete_task_route(owner, task_id) - def session_exists(self, session_id: str) -> bool: + def session_exists(self, owner: str, context_id: str) -> bool: """Return whether Consul still has a live route for a session.""" try: - self._route_for(session_id) + self._route_for(owner, context_id) except ReconnectRequired: return False return True - def task_session_id(self, task_id: str) -> str | None: - """Return the worker session recorded for a task, if any.""" + def task_context_id(self, owner: str, task_id: str) -> str | None: + """Return the owner-scoped context recorded for a task, if any.""" response = requests.get( - f"{self.settings.consul_url}/v1/kv/{_TASK_PREFIX}" + f"{self.settings.consul_url}/v1/kv/{_TASK_PREFIX}{_owner_key(owner)}/" f"{quote(task_id, safe='')}", timeout=10, ) @@ -180,14 +199,15 @@ def task_session_id(self, task_id: str) -> str | None: def _dispatch( self, - session_id: str, + owner: str, + context_id: str, method: A2AMethod | str, message, ) -> WorkerResult: """Send one protobuf-serialized A2A operation to a worker session.""" - route = self._route_for(session_id) + route = self._route_for(owner, context_id) response = requests.post( - f"{route.endpoint}/sessions/{quote(session_id, safe='')}/a2a", + f"{route.endpoint}/sessions/{quote(route.session_id, safe='')}/a2a", data=message.SerializeToString(), headers={ "content-type": PROTOBUF_CONTENT_TYPE, @@ -197,10 +217,10 @@ def _dispatch( **getattr(self, "_worker_request_kwargs", {}), ) if response.status_code in (404, 502): - self._close_worker_session(route, session_id) + self._close_worker_session(route, context_id, owner) raise ReconnectRequired( "Database session ended; reconnect required.", - session_id, + context_id, ) response.raise_for_status() return WorkerResult( @@ -213,13 +233,19 @@ def _dispatch( payload=response.content, ) - def _save_task_route(self, task_id: str, session_id: str) -> None: + def _save_task_route( + self, + owner: str, + task_id: str, + context_id: str, + ) -> None: """Save only task-to-session routing metadata in Consul.""" try: response = requests.put( f"{self.settings.consul_url}/v1/kv/{_TASK_PREFIX}" + f"{_owner_key(owner)}/" f"{quote(task_id, safe='')}", - data=session_id, + data=context_id, timeout=10, ) response.raise_for_status() @@ -228,21 +254,22 @@ def _save_task_route(self, task_id: str, session_id: str) -> None: # are required because workers may use different databases. LOGGER.warning("Could not save route for task %s", task_id) - def _delete_task_route(self, task_id: str) -> None: + def _delete_task_route(self, owner: str, task_id: str) -> None: requests.delete( f"{self.settings.consul_url}/v1/kv/{_TASK_PREFIX}" + f"{_owner_key(owner)}/" f"{quote(task_id, safe='')}", timeout=10, ) - def close_session(self, session_id: str) -> None: + def close_session(self, owner: str, context_id: str) -> None: """Close the child process and remove the Consul route.""" try: - route = self._route_for(session_id) + route = self._route_for(owner, context_id) except ReconnectRequired: - self._delete_route(session_id) + self._delete_route(owner, context_id) return - self._close_worker_session(route, session_id) + self._close_worker_session(route, context_id, owner) def _select_worker(self) -> str: response = requests.get( @@ -279,52 +306,58 @@ def _select_worker(self) -> str: address = service.get("Address") or worker["Node"]["Address"] return f"http://{address}:{service['Port']}" - def _route_for(self, session_id: str) -> SessionRoute: + def _route_for(self, owner: str, context_id: str) -> SessionRoute: response = requests.get( f"{self.settings.consul_url}/v1/kv/{_SESSION_PREFIX}" - f"{quote(session_id, safe='')}", + f"{_owner_key(owner)}/{quote(context_id, safe='')}", timeout=10, ) if response.status_code == 404: raise ReconnectRequired( "Database session expired; reconnect required.", - session_id, + context_id, ) 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) + self._close_worker_session(route, context_id, owner) raise ReconnectRequired( "Database session expired; reconnect required.", - session_id, + context_id, ) return route - def _save_route(self, session_id: str, route: SessionRoute) -> bool: + def _save_route( + self, + owner: str, + context_id: str, + route: SessionRoute, + ) -> bool: response = requests.put( f"{self.settings.consul_url}/v1/kv/{_SESSION_PREFIX}" - f"{quote(session_id, safe='')}?cas=0", + f"{_owner_key(owner)}/{quote(context_id, safe='')}?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: + def _delete_route(self, owner: str, context_id: str) -> None: requests.delete( f"{self.settings.consul_url}/v1/kv/{_SESSION_PREFIX}" - f"{quote(session_id, safe='')}", + f"{_owner_key(owner)}/{quote(context_id, safe='')}", timeout=10, ) def _close_worker_session( self, route: SessionRoute, - session_id: str, + context_id: str, + owner: str | None = None, ) -> None: try: response = requests.delete( - f"{route.endpoint}/sessions/{quote(session_id, safe='')}", + f"{route.endpoint}/sessions/{quote(route.session_id, safe='')}", timeout=10, **getattr(self, "_worker_request_kwargs", {}), ) @@ -333,4 +366,10 @@ def _close_worker_session( except requests.RequestException: pass finally: - self._delete_route(session_id) + if owner is not None: + self._delete_route(owner, context_id) + + +def _owner_key(owner: str) -> str: + """Return a stable non-PII namespace for one authenticated owner.""" + return hashlib.sha256(owner.encode("utf-8")).hexdigest() diff --git a/src/select_ai/cli/a2a.py b/src/select_ai/cli/a2a.py index cab69e8..8e0d622 100644 --- a/src/select_ai/cli/a2a.py +++ b/src/select_ai/cli/a2a.py @@ -5,8 +5,9 @@ # http://oss.oracle.com/licenses/upl. # ----------------------------------------------------------------------------- -import getpass import json +import os +import socket import ssl import click @@ -29,11 +30,30 @@ def a2a(): @a2a.command() -@click.option("--team", "team_name", required=True, help="Database AI team.") +@click.option( + "--deployment", + type=click.Choice(("standalone", "clustered"), case_sensitive=False), + default="standalone", + show_default=True, + help="A2A runtime deployment topology.", +) +@click.option( + "--team", + "team_name", + envvar="SELECT_AI_A2A_TEAM", + help="Database AI team; otherwise collected by the connection form.", +) @click.option("--host", default="127.0.0.1", show_default=True) -@click.option("--port", default=8000, show_default=True, type=int) +@click.option( + "--port", + default=8000, + show_default=True, + envvar="PORT", + type=click.IntRange(min=1, max=65_535), +) @click.option( "--public-url", + envvar="PUBLIC_URL", help="Public base URL advertised in the A2A Agent Card.", ) @click.option("--description", help="A2A agent description.") @@ -44,48 +64,170 @@ def a2a(): type=click.IntRange(min=1), help="Maximum asynchronous Oracle connections.", ) +@click.option( + "--consul-url", + default="http://consul:8500", + show_default=True, + envvar="CONSUL_HTTP_URL", + help="Consul HTTP API URL for clustered deployment.", +) +@click.option( + "--worker-service", + default="select-ai-a2a-worker", + show_default=True, + envvar="WORKER_SERVICE", + help="Consul worker service for clustered deployment.", +) +@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 clustered workers.", +) +@click.option( + "--worker-tls-cert-file", + envvar="WORKER_TLS_CERT_FILE", + type=click.Path(exists=True, dir_okay=False, readable=True), + help="Client certificate presented to clustered workers.", +) +@click.option( + "--worker-tls-key-file", + envvar="WORKER_TLS_KEY_FILE", + type=click.Path(exists=True, dir_okay=False, readable=True), + help="Client private key presented to clustered workers.", +) +@click.option( + "--a2ui-form", + type=click.Path(exists=True, dir_okay=False, readable=True), + help="Custom A2UI connection-form JSON file.", +) +@click.option( + "--require-oauth", + is_flag=True, + envvar="SELECT_AI_A2A_REQUIRE_OAUTH", + help=( + "Require an end-user OAuth bearer token. Without this option, " + "sessions are separated only by the A2A conversation." + ), +) @connection_options def serve( + deployment, team_name, host, port, public_url, description, pool_max_size, + consul_url, + worker_service, + session_ttl_seconds, + worker_tls_ca_file, + worker_tls_cert_file, + worker_tls_key_file, + a2ui_form, + require_oauth, user, password, dsn, wallet_location, wallet_password, ): - """Start an A2A HTTP server for one database AI agent team.""" - if create_app is None or uvicorn is None: + """Start the public A2A server in standalone or clustered mode.""" + if uvicorn is None: raise click.ClickException( "A2A server support requires the optional 'cli' extra. " "Install it with: pip install 'select_ai[cli]'" ) - if password is None: - password = getpass.getpass("Database password: ") - if user is None or dsn is None: - raise click.ClickException( - "--user and --dsn (or their SELECT_AI_* environment variables) " - "are required" - ) if public_url is None: public_url = f"http://{host}:{port}" - app = create_app( - team_name=team_name, - public_url=public_url, - user=user, - password=password, + try: + from select_ai.agent.a2a import ( + ConnectionConfig, + GatewaySettings, + StandaloneSessionSettings, + create_embedded_session_app, + create_gateway_app, + ) + from select_ai.agent.a2a.forms import load_connection_form + except ImportError as error: + raise click.ClickException( + "A2A support requires the optional 'a2a' extra. " + "Install it with: pip install 'select_ai[a2a]'" + ) from error + + connection = ConnectionConfig( dsn=dsn, - wallet_location=wallet_location, - wallet_password=wallet_password, - description=description, - pool_max_size=pool_max_size, + username=user, + password=password, + team_name=team_name, ) + form_template = None + if a2ui_form: + try: + form_template = load_connection_form( + a2ui_form, + connection.missing_fields, + ) + except ValueError as error: + raise click.ClickException(str(error)) from error + + if deployment == "standalone" and not connection.missing_fields: + if create_app is None: + raise click.ClickException( + "Standalone A2A support requires the optional 'cli' extra. " + "Install it with: pip install 'select_ai[cli]'" + ) + app = create_app( + team_name=team_name, + public_url=public_url, + user=user, + password=password, + dsn=dsn, + wallet_location=wallet_location, + wallet_password=wallet_password, + description=description, + pool_max_size=pool_max_size, + require_oauth=require_oauth, + ) + elif deployment == "standalone": + if os.environ.get("WEB_CONCURRENCY", "1") != "1": + raise click.ClickException( + "Dynamic standalone deployment requires WEB_CONCURRENCY=1." + ) + settings = StandaloneSessionSettings( + public_url=public_url, + session_ttl_seconds=session_ttl_seconds, + description=description, + connection=connection, + connection_form_template=form_template, + require_oauth=require_oauth, + ) + app = create_embedded_session_app(settings) + else: + settings = GatewaySettings( + public_url=public_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, + description=description, + connection=connection, + connection_form_template=form_template, + require_oauth=require_oauth, + ) + app = create_gateway_app(settings) click.echo( f"A2A Agent Card: {public_url.rstrip('/')}/.well-known/agent-card.json" @@ -95,7 +237,31 @@ def serve( @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( + "--port", + default=8080, + show_default=True, + type=click.IntRange(min=1, max=65_535), +) +@click.option( + "--worker-id", + envvar="WORKER_ID", + default=socket.gethostname, + show_default="host name", + help="Unique worker ID registered with Consul.", +) +@click.option( + "--consul-url", + envvar="CONSUL_HTTP_URL", + default="http://consul:8500", + show_default=True, + help="Consul HTTP API URL.", +) +@click.option( + "--worker-endpoint", + envvar="WORKER_ENDPOINT", + help="Worker URL advertised through Consul, including scheme and port.", +) @click.option( "--session-ttl-seconds", default=900, @@ -126,6 +292,9 @@ def serve( def worker( host, port, + worker_id, + consul_url, + worker_endpoint, session_ttl_seconds, session_start_timeout_seconds, tls_cert_file, @@ -134,17 +303,23 @@ def worker( ): """Start the internal Select AI session worker.""" try: - from select_ai.agent.a2a import create_worker_app + from select_ai.agent.a2a import WorkerSettings, 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( + settings = WorkerSettings( + consul_url=consul_url, + worker_id=worker_id, + worker_address=os.environ.get("WORKER_ADDRESS", socket.gethostname()), + worker_port=port, session_ttl_seconds=session_ttl_seconds, session_start_timeout_seconds=session_start_timeout_seconds, + worker_endpoint=worker_endpoint, ) + app = create_worker_app(settings) tls_files = (tls_cert_file, tls_key_file, tls_ca_file) if any(tls_files) and not all(tls_files): raise click.ClickException( @@ -162,87 +337,6 @@ def worker( 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 edd594c..199f3e4 100644 --- a/tests/a2a/test_agent_card.py +++ b/tests/a2a/test_agent_card.py @@ -75,9 +75,9 @@ def test_discovery_route_serves_only_the_v03_agent_card(): def test_gateway_card_advertises_a2ui_input_and_output(): app = create_gateway_app( GatewaySettings( - agent_url="https://agent.example.com", + public_url="https://agent.example.com", consul_url="http://consul:8500", - worker_service="select-ai-worker", + worker_service="select-ai-a2a-worker", session_ttl_seconds=900, ) ) diff --git a/tests/a2a/test_worker_runtime.py b/tests/a2a/test_worker_runtime.py index 89671f6..4c26b79 100644 --- a/tests/a2a/test_worker_runtime.py +++ b/tests/a2a/test_worker_runtime.py @@ -33,7 +33,12 @@ TaskNotFoundError, UnsupportedOperationError, ) -from select_ai.agent.a2a import GatewaySettings, session_runtime, worker +from select_ai.agent.a2a import ( + GatewaySettings, + session_process, + session_runtime, + worker, +) from select_ai.agent.a2a.a2ui import ( A2UI_CATALOG_ID, A2UI_EXTENSION_URI, @@ -119,14 +124,19 @@ def process_factory(**kwargs): return process monkeypatch.setattr( - worker.multiprocessing, + session_process.multiprocessing, "Pipe", lambda: (parent, child), ) - monkeypatch.setattr(worker.multiprocessing, "Process", process_factory) - session_worker = worker.SessionWorker(60, 1) - request = worker.OpenSessionRequest( + monkeypatch.setattr( + session_process.multiprocessing, + "Process", + process_factory, + ) + session_worker = session_process.ProcessSessionBackend(60, 1) + request = session_process.SessionSpec( session_id="session-1", + owner="owner-1", username="user", password="password", dsn="database", @@ -145,7 +155,7 @@ def process_factory(**kwargs): assert result == WorkerResult(ResultKind.NONE) assert child.closed - assert processes[0].target is worker._session_process_main + assert processes[0].target is session_process._session_process_main assert processes[0].daemon is True assert parent.sent == [ { @@ -167,14 +177,19 @@ def process_factory(**kwargs): return process monkeypatch.setattr( - worker.multiprocessing, + session_process.multiprocessing, "Pipe", lambda: (parent, child), ) - monkeypatch.setattr(worker.multiprocessing, "Process", process_factory) - session_worker = worker.SessionWorker(60, 1) - request = worker.OpenSessionRequest( + monkeypatch.setattr( + session_process.multiprocessing, + "Process", + process_factory, + ) + session_worker = session_process.ProcessSessionBackend(60, 1) + request = session_process.SessionSpec( session_id="session-1", + owner="owner-1", username="user", password="password", dsn="database", @@ -183,7 +198,10 @@ def process_factory(**kwargs): asyncio.run(session_worker.open(request)) session_worker.sessions["session-1"].expires_at = 0 - with pytest.raises(worker.HTTPException, match="reconnect required"): + with pytest.raises( + session_process.SessionNotFound, + match="reconnect required", + ): asyncio.run(session_worker.get("session-1")) assert process is not None @@ -203,15 +221,24 @@ def process_factory(**kwargs): return process monkeypatch.setattr( - worker.multiprocessing, + session_process.multiprocessing, "Pipe", lambda: (parent, child), ) - monkeypatch.setattr(worker.multiprocessing, "Process", process_factory) - monkeypatch.setattr(worker, "_SESSION_REAPER_INTERVAL_SECONDS", 0.01) - session_worker = worker.SessionWorker(60, 1) - request = worker.OpenSessionRequest( + monkeypatch.setattr( + session_process.multiprocessing, + "Process", + process_factory, + ) + monkeypatch.setattr( + session_process, + "_SESSION_REAPER_INTERVAL_SECONDS", + 0.01, + ) + session_worker = session_process.ProcessSessionBackend(60, 1) + request = session_process.SessionSpec( session_id="session-1", + owner="owner-1", username="user", password="password", dsn="database", @@ -257,8 +284,9 @@ def test_session_runtime_uses_one_async_connection_and_dispatches_a2a( connection_arguments = {} class Runtime: - def __init__(self, session_id, team_name): + def __init__(self, session_id, owner, team_name): assert session_id == "session-1" + assert owner == "owner-1" assert team_name == "TEAM" async def initialize(self): @@ -281,10 +309,10 @@ async def disconnect(): 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(worker, "SessionRuntime", Runtime) + monkeypatch.setattr(session_process, "SessionRuntime", Runtime) asyncio.run( - worker._run_session_process( + session_process._run_session_process( connection, { "user": "user", @@ -292,6 +320,7 @@ async def disconnect(): "dsn": "database", }, "session-1", + "owner-1", "TEAM", ) ) @@ -340,7 +369,7 @@ def __init__(self, **kwargs): lambda *args: "agent-card", ) - runtime = session_runtime.SessionRuntime("session-1", "TEAM") + runtime = session_runtime.SessionRuntime("session-1", "owner-1", "TEAM") asyncio.run(runtime.initialize()) assert initialized == [ @@ -371,7 +400,9 @@ async def initialize(self): monkeypatch.setattr(session_runtime, "OracleTaskStore", Store) monkeypatch.setattr(session_runtime, "OracleContextStore", Store) - runtime = session_runtime.SessionRuntime("session-1", "MISSPELLED_TEAM") + runtime = session_runtime.SessionRuntime( + "session-1", "owner-1", "MISSPELLED_TEAM" + ) with pytest.raises(RuntimeError, match="team does not exist"): asyncio.run(runtime.initialize()) @@ -384,7 +415,7 @@ class Handler: async def on_get_task(self, _request, _context): raise TaskNotFoundError - runtime = session_runtime.SessionRuntime("session-1", "TEAM") + runtime = session_runtime.SessionRuntime("session-1", "owner-1", "TEAM") runtime.handler = Handler() result = asyncio.run( @@ -462,7 +493,11 @@ def post(*args, **kwargs): monkeypatch.setattr( client, "_route_for", - lambda session_id: type("Route", (), {"endpoint": "http://worker"})(), + lambda owner, context_id: type( + "Route", + (), + {"endpoint": "http://worker", "session_id": "session-1"}, + )(), ) request = SendMessageRequest( @@ -472,7 +507,7 @@ def post(*args, **kwargs): parts=[new_text_part("hello")], ) ) - result = client.send_message("context-1", request) + result = client.send_message("owner-1", "context-1", request) assert result is not None assert result.message_id == "m1" @@ -486,9 +521,9 @@ def post(*args, **kwargs): def test_worker_client_uses_consul_https_endpoint_with_mtls(monkeypatch): settings = GatewaySettings( - agent_url="https://gateway.example.com", + public_url="https://gateway.example.com", consul_url="http://consul:8500", - worker_service="select-ai-worker", + worker_service="select-ai-a2a-worker", session_ttl_seconds=60, worker_tls_ca_file="/tls/ca.pem", worker_tls_cert_file="/tls/gateway.pem", @@ -527,9 +562,9 @@ def json(): def test_mtls_requires_all_three_gateway_files(): with pytest.raises(ValueError, match="worker mTLS requires"): GatewaySettings( - agent_url="https://gateway.example.com", + public_url="https://gateway.example.com", consul_url="http://consul:8500", - worker_service="select-ai-worker", + worker_service="select-ai-a2a-worker", session_ttl_seconds=60, worker_tls_ca_file="/tls/ca.pem", ) @@ -569,6 +604,7 @@ async def put(self, _url, json): ) assert registered["Meta"] == {"endpoint": "https://worker-0.internal"} + assert registered["Name"] == "select-ai-a2a-worker" def test_a2ui_operation_uses_a_metadata_marked_data_part(): @@ -642,7 +678,7 @@ def test_connection_action_accepts_gemini_unmarked_data_part(): "version": "v0.9", "action": { "name": "submit_database_connection", - "context": {"team_name": "TEAM"}, + "context": {"ai_agent": "TEAM"}, }, } ) @@ -653,7 +689,7 @@ def test_connection_action_accepts_gemini_unmarked_data_part(): assert action == { "name": "submit_database_connection", - "context": {"team_name": "TEAM"}, + "context": {"ai_agent": "TEAM"}, } @@ -662,7 +698,7 @@ def test_gateway_returns_connection_error_when_worker_rejects_opening(): class Client: @staticmethod - def open_session(_context_id, _session_info): + def open_session(_owner, _context_id, _session_info): raise requests.HTTPError("worker rejected the connection") handler.worker_client = Client() @@ -670,11 +706,12 @@ def open_session(_context_id, _session_info): session_id = asyncio.run( handler._open_session( { - "dsn": "database", + "connection_url": "database", "username": "user", "password": "password", - "team_name": "TEAM", + "ai_agent": "TEAM", }, + "owner-1", "context-1", ) ) @@ -688,14 +725,15 @@ def __init__(self): self.opened = False self.saved = None - def session_exists(self, _session_id): + def session_exists(self, _owner, _context_id): return self.opened - def open_session(self, session_id, session_info): - assert session_id == "context-1" + def open_session(self, owner, context_id, session_info): + assert owner == "a2a-conversation" + assert context_id == "context-1" assert session_info.team_name == "TEAM" self.opened = True - return session_id + return "session-1" client = Client() handler = GatewayRequestHandler.__new__(GatewayRequestHandler) @@ -731,10 +769,10 @@ def open_session(self, session_id, session_info): "action": { "name": "submit_database_connection", "context": { - "dsn": "database", + "connection_url": "database", "username": "user", "password": "password", - "team_name": "TEAM", + "ai_agent": "TEAM", }, }, } @@ -756,11 +794,11 @@ class Client: calls = 0 @staticmethod - def session_exists(_session_id): + def session_exists(_owner, _context_id): return True - def send_message(self, session_id, request): - assert session_id == "context-1" + def send_message(self, _owner, context_id, request): + assert context_id == "context-1" self.calls += 1 if self.calls == 1: assert request.message.task_id == ( @@ -770,7 +808,7 @@ def send_message(self, session_id, request): assert not request.message.task_id return Task( id="worker-task", - context_id=session_id, + context_id=context_id, status={"state": TaskState.TASK_STATE_COMPLETED}, ) @@ -797,11 +835,11 @@ def send_message(self, session_id, request): def test_gateway_preserves_missing_worker_task_error(): class Client: @staticmethod - def session_exists(_session_id): + def session_exists(_owner, _context_id): return True @staticmethod - def send_message(_session_id, _request): + def send_message(_owner, _context_id, _request): raise TaskNotFoundError handler = GatewayRequestHandler.__new__(GatewayRequestHandler) @@ -823,8 +861,8 @@ def send_message(_session_id, _request): def test_gateway_bootstrap_without_context_returns_transient_form_task(): class Client: @staticmethod - def session_exists(session_id): - assert session_id + def session_exists(_owner, context_id): + assert context_id return False handler = GatewayRequestHandler.__new__(GatewayRequestHandler) @@ -856,7 +894,7 @@ def session_exists(session_id): def test_gateway_returns_task_form_when_task_session_expires_before_send(): class Client: @staticmethod - def get_task(_request): + def get_task(_owner, _request): raise ReconnectRequired( "Database session ended; reconnect required.", "context-1", @@ -886,7 +924,7 @@ def get_task(_request): def test_gateway_returns_task_form_when_task_route_is_missing_before_send(): class Client: @staticmethod - def get_task(_request): + def get_task(_owner, _request): return None handler = GatewayRequestHandler.__new__(GatewayRequestHandler) @@ -913,11 +951,11 @@ def get_task(_request): def test_gateway_uses_task_form_for_expired_context_only_message(): class Client: @staticmethod - def session_exists(_session_id): + def session_exists(_owner, _context_id): return True @staticmethod - def send_message(_session_id, _request): + def send_message(_owner, _context_id, _request): raise ReconnectRequired( "Database session ended; reconnect required.", "context-1", @@ -947,7 +985,7 @@ def send_message(_session_id, _request): def test_gateway_uses_connection_form_when_get_task_session_expires(): class Client: @staticmethod - def get_task(_request): + def get_task(_owner, _request): raise ReconnectRequired( "Database session ended; reconnect required.", "context-1", @@ -973,7 +1011,7 @@ def get_task(_request): def test_gateway_uses_connection_form_when_get_task_route_is_missing(): class Client: @staticmethod - def get_task(_request): + def get_task(_owner, _request): return None handler = GatewayRequestHandler.__new__(GatewayRequestHandler) @@ -996,7 +1034,7 @@ def get_task(_request): def test_gateway_uses_connection_form_when_cancel_task_session_expires(): class Client: @staticmethod - def cancel_task(_task_id): + def cancel_task(_owner, _task_id): raise ReconnectRequired( "Database session ended; reconnect required.", "context-1", @@ -1022,7 +1060,7 @@ def cancel_task(_task_id): def test_gateway_uses_connection_form_when_cancel_task_route_is_missing(): class Client: @staticmethod - def cancel_task(_task_id): + def cancel_task(_owner, _task_id): return None handler = GatewayRequestHandler.__new__(GatewayRequestHandler) @@ -1045,7 +1083,7 @@ def cancel_task(_task_id): def test_gateway_exposes_connection_form_data_when_list_session_expires(): class Client: @staticmethod - def list_tasks(_context_id, _request): + def list_tasks(_owner, _context_id, _request): raise ReconnectRequired( "Database session ended; reconnect required.", "context-1",