Refactor A2A gateway routing and session lifecycle

  - simplify gateway request handling and A2UI bootstrap flow
  - add database-backed task routing and worker protocol types
  - clean up expired worker child processes
  - improve GCP gateway architecture and mTLS deployment
  - expand A2A gateway and worker tests

Author: aosingh

gcloud/gateway/README.md

@@ -1,20 +1,67 @@
-# Dynamic gateway deployment
+# Dynamic gateway
 
 Dynamic gateway mode exposes one public A2A endpoint. Each user dynamically
 selects an Oracle database connection and Select AI team through the A2UI
 connection form. A session remains available for 15 minutes by default. Set a
 different lifetime in seconds with `--session-ttl-seconds`; for example,
 `--session-ttl-seconds 1800` keeps sessions for 30 minutes.
 
+## Protocol architecture
+
+```text
+┌──────────────┐  A2A JSON-RPC/HTTP  ┌──────────────────┐                                     ┌────────────────────────────┐
+│ A2A client   │────────────────────►│ Gateway instances│──── route lookup/update ──────────► │ Service Registry           │
+└──────────────┘                     │ Public A2A API   │                                     │ Service discovery          │
+                                     │ A2UI bootstrap   │                                     │ Session routes             │
+                                     └────────┬─────────┘                                     │ Task routes                │
+                                              │ Internal protobuf                             │                            │
+                                              ▼                                               │                            │
+                                     ┌──────────────────┐                                     │                            │
+                                     │ Worker pool      │──── registration / heartbeat------->│                            │
+                                     │ worker-0, ...    │                                     │                            │
+                                     └────────┬─────────┘                                     └────────────────────────────┘
+                                              │ one child per database session
+                                              ▼
+                                     ┌──────────────────┐
+                                     │ Session runtime  │
+                                     │ A2A handler      │
+                                     │ Task/context     │
+                                     │ stores           │
+                                     │ Database session │
+                                     └────────┬─────────┘
+                                              │ SQL / Select AI
+                                              ▼
+                                     ┌──────────────────┐
+                                     │ Oracle Database  │
+                                     └──────────────────┘
+```
+
+The gateway 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,
+`DefaultRequestHandler`, `OracleTaskStore`, and `OracleContextStore`.
+
+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
+or routed.
+
+## GCP deployment
+
+The protocol architecture above is implemented on GCP as follows:
+
 ```text
  ┌──────────────────────┐
  │ A2A / Gemini client  │
  └──────────┬───────────┘
             │ public A2A
             v
- ┌──────────────────────┐
- │ Cloud Run gateway    │
- └──────┬───────┬───────┘
+ ┌────────────────────────────┐
+ │ Cloud Run gateway          │
+ │ A2A proxy + form bootstrap │
+ └──────┬───────────┬─────────┘
         │       │ private VPC: mTLS request to worker hostname
         │       │
         │       │     ┌─────────────────────── GKE ───────────────────────┐
@@ -23,19 +70,16 @@ different lifetime in seconds with `--session-ttl-seconds`; for example,
         │             │                    │                              │
         │             │                    v                              │
         │             │ [StatefulSet worker-0 / worker-1 / ...]           │
-        │             │ session child process → Oracle Database           │
+        │             │ session child: A2A handler + Oracle stores        │
+        │             │                    │                              │
+        │             │                    v                              │
+        │             │              Oracle Database                      │
         │             │                                                   │
         │             │ [Consul]                                          │
         └────────────>│ selects healthy worker; returns worker hostname   │
                       └───────────────────────────────────────────────────┘
 ```
 
-The gateway is the only public A2A application. Consul and workers are a GKE
-clustered service: Consul selects a worker for each new dynamic session, and
-the chosen worker retains that session's process and Oracle conversation.
-
-## Deploy the complete stack
-
 Run this from the repository root:
 
 ```bash
@@ -68,10 +112,14 @@ gcloud/gateway/deploy.sh \
 ```
 
 The Cloud Run gateway uses direct VPC egress to reach the internal Consul load
-balancer and GKE worker pod addresses. The default one-instance gateway limit
-is intentional: gateway A2A task and context/session state is currently in
-memory. Workers, rather than the gateway, provide the clustered capacity for
-dynamic sessions.
+balancer and GKE worker pod addresses. The gateway 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
+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

gcloud/gateway/deploy.sh

@@ -166,10 +166,15 @@ if [[ "$enable_worker_mtls" == "true" ]]; then
     fi
   done
   if [[ "$create_mtls_material" != "true" ]]; then
-    existing_worker_certificate_sans="$(gcloud secrets versions access latest \
-      --secret=select-ai-worker-mtls-cert --project="$project_id" 2>/dev/null | \
-      openssl x509 -noout -ext subjectAltName 2>/dev/null || true)"
-    if [[ "$existing_worker_certificate_sans" != *"DNS:$worker_certificate_dns_name"* ]]; then
+    # macOS ships LibreSSL, which does not support x509's -ext option.
+    # 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 | \
+      openssl x509 -noout -text 2>/dev/null || true
+    )"
+    if ! printf '%s\n' "$existing_worker_certificate_text" | \
+      grep -F -- "DNS:$worker_certificate_dns_name" >/dev/null; then
       create_mtls_material="true"
       echo "Replacing mTLS material because the worker certificate does not"
       echo "match the GKE DNS domain."

src/select_ai/agent/a2a/a2ui.py

@@ -5,24 +5,34 @@
 # https://oss.oracle.com/licenses/upl.
 # -----------------------------------------------------------------------------
 
-"""Shared A2UI protocol declarations for Select AI A2A agents."""
+"""Shared A2UI protocol helpers for Select AI A2A agents."""
 
+from __future__ import annotations
+
+from collections.abc import Iterator
+
+from a2a.helpers import new_data_part
 from a2a.types import AgentExtension
-from google.protobuf.json_format import ParseDict
+from a2a.types.a2a_pb2 import Message, Part
+from google.protobuf.json_format import MessageToDict, ParseDict
 from google.protobuf.struct_pb2 import Struct
 
-from select_ai.agent.a2a.forms import _CATALOG
-
-A2UI_EXTENSION_URI = "https://a2ui.org/a2a-extension/a2ui/v0.9"
+A2UI_VERSION = "v0.9"
+A2UI_EXTENSION_URI = f"https://a2ui.org/a2a-extension/a2ui/{A2UI_VERSION}"
 A2UI_MIME_TYPE = "application/json+a2ui"
+A2UI_CATALOG_ID = (
+    "https://www.gstatic.com/vertexaisearch/a2ui/"
+    f"{A2UI_VERSION.replace('.', '_')}/"
+    "gemini_enterprise_composite_catalog.json"
+)
 
 
 def a2ui_extension() -> AgentExtension:
-    """Return the A2UI v0.9 capability used by Gemini Enterprise."""
+    """Return the supported A2UI capability used by Gemini Enterprise."""
     params = ParseDict(
         {
             "acceptsInlineCatalogs": True,
-            "supportedCatalogIds": [_CATALOG],
+            "supportedCatalogIds": [A2UI_CATALOG_ID],
         },
         Struct(),
     )
@@ -31,3 +41,57 @@ def a2ui_extension() -> AgentExtension:
         description="Provides agent driven UI using the A2UI JSON format.",
         params=params,
     )
+
+
+def a2ui_part(operation: dict) -> Part:
+    """Encode one A2UI operation in an A2A data part."""
+    part = new_data_part(operation)
+    ParseDict({"mimeType": A2UI_MIME_TYPE}, part.metadata)
+    return part
+
+
+def a2ui_operations(message: Message) -> Iterator[dict]:
+    """Yield operations from data parts marked with the A2UI MIME type."""
+    for part in message.parts:
+        if not _is_a2ui_part(part):
+            continue
+        data = MessageToDict(part.data)
+        operations = data if isinstance(data, list) else [data]
+        yield from (
+            operation
+            for operation in operations
+            if isinstance(operation, dict)
+        )
+
+
+def find_action(message: Message, name: str) -> dict | None:
+    """Return a named A2UI action, including Gemini's unmarked input form."""
+    for part in message.parts:
+        if part.WhichOneof("content") != "data":
+            continue
+        mime_type = MessageToDict(part.metadata).get("mimeType")
+        # Gemini Enterprise does not currently echo the A2UI MIME metadata on
+        # a submitted form action.  Accept that legacy input shape, while
+        # still rejecting data parts explicitly marked as another format.
+        if mime_type not in (None, A2UI_MIME_TYPE):
+            continue
+        data = MessageToDict(part.data)
+        operations = data if isinstance(data, list) else [data]
+        for operation in operations:
+            if (
+                not isinstance(operation, dict)
+                or operation.get("version") != A2UI_VERSION
+            ):
+                continue
+            action = operation.get("action")
+            if isinstance(action, dict) and action.get("name") == name:
+                return action
+    return None
+
+
+def _is_a2ui_part(part: Part) -> bool:
+    """Identify an A2UI data part by its standard MIME type."""
+    return (
+        part.WhichOneof("content") == "data"
+        and MessageToDict(part.metadata).get("mimeType") == A2UI_MIME_TYPE
+    )

src/select_ai/agent/a2a/forms.py

@@ -7,27 +7,30 @@
 
 """A2UI connection form emitted by the public gateway."""
 
+from uuid import uuid4
 
-_CATALOG = (
-    "https://www.gstatic.com/vertexaisearch/a2ui/v0_9/"
-    "gemini_enterprise_composite_catalog.json"
-)
+from select_ai.agent.a2a.a2ui import A2UI_CATALOG_ID, A2UI_VERSION
 
 
-def connection_form() -> list[dict]:
+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.
+    surface_id = surface_id or f"db-connect-{uuid4().hex}"
     return [
         {
-            "version": "v0.9",
+            "version": A2UI_VERSION,
             "createSurface": {
-                "surfaceId": "db-connect",
-                "catalogId": _CATALOG,
+                "surfaceId": surface_id,
+                "catalogId": A2UI_CATALOG_ID,
             },
         },
         {
-            "version": "v0.9",
+            "version": A2UI_VERSION,
             "updateComponents": {
-                "surfaceId": "db-connect",
+                "surfaceId": surface_id,
                 "components": [
                     {"id": "root", "component": "Card", "child": "column"},
                     {
@@ -102,9 +105,9 @@ def connection_form() -> list[dict]:
             },
         },
         {
-            "version": "v0.9",
+            "version": A2UI_VERSION,
             "updateDataModel": {
-                "surfaceId": "db-connect",
+                "surfaceId": surface_id,
                 "path": "/",
                 "value": {
                     "dsn": "",