From b154baf6d72979e8bf4fbea42fb3333fd61b08fc Mon Sep 17 00:00:00 2001
From: Jinwoo Hong <73622457+Jinwoo-H@users.noreply.github.com>
Date: Sat, 10 Oct 2026 19:53:44 -0400
Subject: [PATCH] fix(relay): staging director deploys run as relay-dir with
 rehome identity, and can register one cell's cap (#27118)

Rebased onto main as one commit after its parent squash-merged.
---
 ...cloud-bootstrap-relay-staging-capacity.yml |  1 -
 .../workflows/cloud-deploy-relay-staging.yml  | 95 ++++++++++++++++++-
 .../cloud-prove-relay-staging-capacity.yml    |  1 -
 .../relay-deploy-workflow-argv.test.mjs       | 20 ++++
 .../relay-staging-deploy-identity.test.mjs    | 10 ++
 5 files changed, 122 insertions(+), 5 deletions(-)

diff --git a/.github/workflows/cloud-bootstrap-relay-staging-capacity.yml b/.github/workflows/cloud-bootstrap-relay-staging-capacity.yml
index 29e68510b21..4ca3e356792 100644
--- a/.github/workflows/cloud-bootstrap-relay-staging-capacity.yml
+++ b/.github/workflows/cloud-bootstrap-relay-staging-capacity.yml
@@ -544,7 +544,6 @@ jobs:
               --capacity-service-account "${CAPACITY_SERVICE_ACCOUNT}" \
               --capacity-cell-id "${cell_id}" \
               --director-cells-json "${desired_cells_json}" \
-              --min-instances 0 \
               --prune-revisions true \
               --release-id "bootstrap-${cell_id}-${GITHUB_RUN_ID}-${GITHUB_RUN_ATTEMPT}")"
             echo "${director_result}"
diff --git a/.github/workflows/cloud-deploy-relay-staging.yml b/.github/workflows/cloud-deploy-relay-staging.yml
index bd82fbfb851..4e74c835100 100644
--- a/.github/workflows/cloud-deploy-relay-staging.yml
+++ b/.github/workflows/cloud-deploy-relay-staging.yml
@@ -13,6 +13,35 @@ on:
         default: preserve
         type: choice
         options: [preserve, 'off', dry-run, 'on']
+      expected-rehome-generation:
+        description: Exact durable regional-rehome generation; it must remain disabled
+        required: true
+        type: string
+      bootstrap-runtime-identity:
+        description: One-time move from the stamped-cell identity to the director identity
+        required: true
+        default: false
+        type: boolean
+      predecessor-image-digest:
+        description: Bootstrap only. Exact immutable serving predecessor digest
+        required: false
+        type: string
+      capacity-cell-id:
+        description: Register this cell's reviewed staging.tfvars cap and bound with the director
+        required: true
+        default: none
+        type: choice
+        options: [none, staging-gce-c1, staging-gce-c2, staging-gce-c3, staging-gce-c4]
+      min-instances:
+        description: Director floor; preserve keeps the serving revision's, so a reserve run's 2 survives
+        required: true
+        default: preserve
+        type: choice
+        options: [preserve, '0', '2']
+      confirmation:
+        description: BOOTSTRAP_RELAY_DIRECTOR_REHOME_IDENTITY for a bootstrap
+        required: false
+        type: string
       shadow-seat-feed-cells:
         description: Preserve the step-3 shadow seat-feed cells, or set all or a comma list of cell ids; dropping a reserve cell is refused
         required: true
@@ -47,11 +76,34 @@ jobs:
       CAPACITY_SERVICE_ACCOUNT: ${{ vars.STAGING_GCP_RELAY_CAPACITY_SERVICE_ACCOUNT }}
       ASIA_PROOF_SERVICE_ACCOUNT: ${{ vars.STAGING_GCP_RELAY_ASIA_PROOF_SERVICE_ACCOUNT }}
       REGIONAL_PLACEMENT_SECRET: orca-cloud-relay-regional-placement-enabled
+      DIRECTOR_RUNTIME_SERVICE_ACCOUNT: orca-cloud-staging-relay-dir@onorca-cloud-staging.iam.gserviceaccount.com
+      PREDECESSOR_RUNTIME_SERVICE_ACCOUNT: orca-cloud-staging-relay@onorca-cloud-staging.iam.gserviceaccount.com
+      REHOME_AUDIENCE: https://relay-staging.onorca.dev/v1/admin/host-drain
+      EXPECTED_REHOME_GENERATION: ${{ inputs.expected-rehome-generation }}
+      BOOTSTRAP_RUNTIME_IDENTITY: ${{ inputs.bootstrap-runtime-identity }}
+      PREDECESSOR_IMAGE_DIGEST: ${{ inputs.predecessor-image-digest }}
+      CAPACITY_CELL_ID: ${{ inputs.capacity-cell-id }}
+      MIN_INSTANCES: ${{ inputs.min-instances }}
     steps:
       - uses: actions/checkout@v4
 
-      - name: Require the expected immutable image
-        run: '[[ "${EXPECTED_IMAGE_DIGEST}" =~ ^sha256:[a-f0-9]{64}$ ]]'
+      - name: Require the expected immutable image and identity request
+        shell: bash
+        env:
+          CONFIRMATION: ${{ inputs.confirmation }}
+        run: |
+          set -euo pipefail
+          [[ "${EXPECTED_IMAGE_DIGEST}" =~ ^sha256:[a-f0-9]{64}$ ]]
+          [[ "${EXPECTED_REHOME_GENERATION}" =~ ^(0|[1-9][0-9]*)$ ]]
+          [[ "${BOOTSTRAP_RUNTIME_IDENTITY}" =~ ^(true|false)$ ]]
+          [[ "${MIN_INSTANCES}" =~ ^(preserve|0|2)$ ]]
+          if test "${BOOTSTRAP_RUNTIME_IDENTITY}" = true; then
+            [[ "${PREDECESSOR_IMAGE_DIGEST}" =~ ^sha256:[a-f0-9]{64}$ ]]
+            test "${CAPACITY_CELL_ID}" = none
+            test "${CONFIRMATION}" = BOOTSTRAP_RELAY_DIRECTOR_REHOME_IDENTITY
+          else
+            test -z "${PREDECESSOR_IMAGE_DIGEST}"
+          fi
 
       - uses: hashicorp/setup-terraform@v3
         with:
@@ -79,6 +131,16 @@ jobs:
           test "${IMAGE}" = \
             "${GCP_REGION}-docker.pkg.dev/${GCP_PROJECT_ID}/${REPOSITORY_ID}/${IMAGE_NAME}@${EXPECTED_IMAGE_DIGEST}"
           echo "IMAGE=${IMAGE}" >> "${GITHUB_ENV}"
+          if test "${CAPACITY_CELL_ID}" != none; then
+            # The deploy script refuses any topology change beyond this one cell's cap and bound.
+            DESIRED_CELLS_JSON="$(terraform -chdir=infra/terraform console \
+              -var-file=environments/staging.tfvars \
+              <<< 'local.relay_director_cells_json' | jq -er '.')"
+            jq -e --arg cell "${CAPACITY_CELL_ID}" \
+              'any(.[]; .id == $cell and .connectionHardCap != null)' \
+              <<< "${DESIRED_CELLS_JSON}" >/dev/null
+            echo "DESIRED_CELLS_JSON=${DESIRED_CELLS_JSON}" >> "${GITHUB_ENV}"
+          fi
 
       - uses: google-github-actions/setup-gcloud@v2
 
@@ -109,6 +171,26 @@ jobs:
             --format='value(name)' | awk -F/ '{print $NF}')"
           [[ "${regional_version}" =~ ^[1-9][0-9]*$ ]]
           RELEASE_ID="${GITHUB_RUN_ID}-${GITHUB_RUN_ATTEMPT}-${GITHUB_SHA:0:8}"
+          identity_args=()
+          if test "${BOOTSTRAP_RUNTIME_IDENTITY}" = true; then
+            identity_args=(
+              --bootstrap-runtime-identity true
+              --predecessor-runtime-service-account "${PREDECESSOR_RUNTIME_SERVICE_ACCOUNT}"
+              --predecessor-image-digest "${PREDECESSOR_IMAGE_DIGEST}"
+            )
+          fi
+          # Omitted, the deploy script inherits the serving revision's floor.
+          floor_args=()
+          if test "${MIN_INSTANCES}" != preserve; then
+            floor_args=(--min-instances "${MIN_INSTANCES}")
+          fi
+          topology_args=()
+          if test "${CAPACITY_CELL_ID}" != none; then
+            topology_args=(
+              --capacity-cell-id "${CAPACITY_CELL_ID}"
+              --director-cells-json "${DESIRED_CELLS_JSON}"
+            )
+          fi
           node dev/scripts/deploy-relay-blue-green.mjs \
             --project "${GCP_PROJECT_ID}" \
             --region "${GCP_REGION}" \
@@ -119,7 +201,14 @@ jobs:
             --capacity-service-account "${CAPACITY_SERVICE_ACCOUNT}" \
             --asia-proof-service-account "${ASIA_PROOF_SERVICE_ACCOUNT}" \
             --regional-placement-secret-version "${regional_version}" \
-            --min-instances 0 \
+            "${floor_args[@]}" \
+            --runtime-service-account "${DIRECTOR_RUNTIME_SERVICE_ACCOUNT}" \
+            "${identity_args[@]}" \
+            --rehome-director-service-account "${DIRECTOR_RUNTIME_SERVICE_ACCOUNT}" \
+            --rehome-audience "${REHOME_AUDIENCE}" \
+            --rehome-control-origin https://relay-staging.onorca.dev \
+            --expected-rehome-generation "${EXPECTED_REHOME_GENERATION}" \
+            "${topology_args[@]}" \
             --admin-audience https://relay-staging.onorca.dev/v1/admin/drain \
             --reserve-placement "${RESERVE_PLACEMENT}" \
             --shadow-seat-feed-cells "${SHADOW_SEAT_FEED_CELLS}" \
diff --git a/.github/workflows/cloud-prove-relay-staging-capacity.yml b/.github/workflows/cloud-prove-relay-staging-capacity.yml
index 279007cbc14..15d61e91dc6 100644
--- a/.github/workflows/cloud-prove-relay-staging-capacity.yml
+++ b/.github/workflows/cloud-prove-relay-staging-capacity.yml
@@ -417,7 +417,6 @@ jobs:
             --capacity-service-account "${CAPACITY_SERVICE_ACCOUNT}" \
             --capacity-cell-id "${TARGET_CELL_ID}" \
             --director-cells-json "${DESIRED_CELLS_JSON}" \
-            --min-instances 0 \
             --prune-revisions true \
             --release-id "${RELEASE_ID}")"
           echo "${DEPLOY_RESULT}"
diff --git a/cloud/dev/scripts/relay-deploy-workflow-argv.test.mjs b/cloud/dev/scripts/relay-deploy-workflow-argv.test.mjs
index a4861625084..8dea82fe6e9 100644
--- a/cloud/dev/scripts/relay-deploy-workflow-argv.test.mjs
+++ b/cloud/dev/scripts/relay-deploy-workflow-argv.test.mjs
@@ -40,6 +40,19 @@ function sample(flag) {
   return `sample-${flag}`
 }
 
+// Flags a workflow passes through a bash array (`"${name[@]}"`), from every `name=(...)` it
+// assigns; parsing them all at once is the widest set that workflow can send.
+function arrayArguments(source, name) {
+  const flags = []
+  for (const match of source.matchAll(new RegExp(`\\b${name}=\\(([^)]*)\\)`, 'g'))) {
+    for (const entry of match[1].matchAll(/(--[a-z0-9-]+)\s+("[^"]*"|\S+)/g)) {
+      flags.push([entry[1], entry[2].replace(/^"(.*)"$/, '$1')])
+    }
+  }
+  if (flags.length === 0) throw new Error(`no flags assigned to ${name}`)
+  return flags
+}
+
 function invocations(source, script) {
   const lines = source.split('\n')
   const found = []
@@ -49,6 +62,13 @@ function invocations(source, script) {
     let line = lines[index]
     while (line.trimEnd().endsWith('\\')) {
       line = lines[++index]
+      const array = /^\s*"\$\{([a-z_]+)\[@\]\}"\s*\\?\s*$/.exec(line)
+      if (array) {
+        for (const [flag, value] of arrayArguments(source, array[1])) {
+          argv.push(flag, value.includes('$') ? sample(flag.slice(2)) : value)
+        }
+        continue
+      }
       const match = /^\s*(--[a-z0-9-]+)\s+(.+?)\s*\\?\s*$/.exec(line)
       if (!match) throw new Error(`unparsed ${script} argument line: ${line.trim()}`)
       const [, flag, raw] = match
diff --git a/cloud/dev/scripts/relay-staging-deploy-identity.test.mjs b/cloud/dev/scripts/relay-staging-deploy-identity.test.mjs
index 8afbfb5ca94..8b7fb069915 100644
--- a/cloud/dev/scripts/relay-staging-deploy-identity.test.mjs
+++ b/cloud/dev/scripts/relay-staging-deploy-identity.test.mjs
@@ -223,3 +223,13 @@ test('the identity is exposed through its own outputs', () => {
   assert.match(outputs, /output "github_staging_relay_deploy_workload_identity_provider"/)
   assert.match(outputs, /output "github_staging_relay_deploy_service_account"/)
 })
+
+// Why: relay-deploy-workflow-argv.test.mjs proves the argv parses; this pins the identity it names.
+test('the staging director deploy runs as relay-dir with its rehome identity', () => {
+  const source = workflow('deploy-relay-staging.yml')
+  const start = source.indexOf('node dev/scripts/deploy-relay-blue-green.mjs')
+  const block = source.slice(start, source.indexOf('\n\n', start))
+  assert.match(block, /--runtime-service-account "\$\{DIRECTOR_RUNTIME_SERVICE_ACCOUNT\}"/)
+  assert.match(block, /--rehome-director-service-account "\$\{DIRECTOR_RUNTIME_SERVICE_ACCOUNT\}"/)
+  assert.match(source, /DIRECTOR_RUNTIME_SERVICE_ACCOUNT: orca-cloud-staging-relay-dir@/)
+})