From aa3dc15d0cc821a81769da351e57ccec35daba24 Mon Sep 17 00:00:00 2001
From: Yu Hu <yuhuyoyo@verily.health>
Date: Mon, 10 Mar 2025 14:27:47 -0400
Subject: [PATCH 01/11] fix app proxy for non-google

---
 startupscript/dataproc/startup.sh | 27 ++++++++++++++++++---------
 1 file changed, 18 insertions(+), 9 deletions(-)

diff --git a/startupscript/dataproc/startup.sh b/startupscript/dataproc/startup.sh
index 9c3c0ca2..d01f912d 100644
--- a/startupscript/dataproc/startup.sh
+++ b/startupscript/dataproc/startup.sh
@@ -197,21 +197,30 @@ function set_guest_attributes() {
 # Map the CLI server to appropriate service path
 function get_service_url() {
   case "$1" in
-    "verily") echo "https://terra-$2.api.verily.com" ;;
-    "verily-devel") echo "https://terra-devel-$2.api.verily.com" ;;
-    "verily-autopush") echo "https://terra-autopush-$2.api.verily.com" ;;
-    "verily-staging") echo "https://terra-staging-$2.api.verily.com" ;;
-    "verily-preprod") echo "https://terra-preprod-$2.api.verily.com" ;;
+    "verily") echo "https://workbench.verily.com/api/$2" ;;
     "dev-stable") echo "https://workbench-dev.verily.com/api/$2" ;;
     "dev-unstable") echo "https://workbench-dev-unstable.verily.com/api/$2" ;;
     "test") echo "https://workbench-test.verily.com/api/$2" ;;
     "staging") echo "https://workbench-staging.verily.com/api/$2" ;;
-    "prod") echo "https://workbench-prod.verily.com/api/$2" ;;
+    "prod") echo "https://workbench.verily.com/api/$2" ;;
     *) return 1 ;;
   esac
 }
 readonly -f get_service_url
 
+function get_app_proxy_url() {
+  local env=$1
+  case "$env" in
+    verily)
+      env="prod"
+      ;;
+    dev-stable)
+      env="dev"
+      ;;
+  esac
+  echo "https://workbench-app-${env}.verily.com"
+}
+
 # If the script exits without error let the UI know it completed successfully
 # Otherwise if an error occurred write the line and command that failed to guest attributes.
 function exit_handler {
@@ -798,7 +807,7 @@ chown "${LOGIN_USER}:${LOGIN_USER}" "${USER_BASH_PROFILE}"
 
 # TODO(BENCH-2612): use workbench CLI instead to get user profile.
 IS_NON_GOOGLE_ACCOUNT="$(curl "${USER_SERVICE_URL}/api/profile?path=non_google_account" \
-                    -H "accept: application/json" -H "Authorization: Bearer $(gcloud auth print-access-token)" \
+                    -H "accept: application/json" -H "Authorization: Bearer $(wb auth print-access-token)" \
                   | jq '.value')"
 readonly IS_NON_GOOGLE_ACCOUNT
 
@@ -808,12 +817,12 @@ if [[ "${IS_NON_GOOGLE_ACCOUNT}" == "true" ]]; then
 # Start workbench app proxy agent 
 ###################################
 
-  APP_PROXY="$(get_metadata_value "instance/attributes/terra-app-proxy")"
+  APP_PROXY="$(get_app_proxy_url "${TERRA_SERVER}")"
   readonly APP_PROXY
   if [[ -n "${APP_PROXY}" ]]; then
     emit "Using custom Proxy Agent"
     RESOURCE_ID="$(get_metadata_value "instance/attributes/terra-resource-id")"
-    NEW_PROXY="https://${APP_PROXY}"
+    NEW_PROXY="${APP_PROXY}"
     NEW_PROXY_URL="${RESOURCE_ID}.${APP_PROXY}"
     readonly RESOURCE_ID
     readonly NEW_PROXY

From cdd8c53f422c7988ca7b96e499f53ed3e4840bff Mon Sep 17 00:00:00 2001
From: Yu Hu <yuhuyoyo@verily.health>
Date: Mon, 10 Mar 2025 14:30:33 -0400
Subject: [PATCH 02/11] add comment

---
 startupscript/dataproc/startup.sh | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/startupscript/dataproc/startup.sh b/startupscript/dataproc/startup.sh
index d01f912d..3f4386cf 100644
--- a/startupscript/dataproc/startup.sh
+++ b/startupscript/dataproc/startup.sh
@@ -208,18 +208,20 @@ function get_service_url() {
 }
 readonly -f get_service_url
 
+# Get the app proxy url for the env.
 function get_app_proxy_url() {
-  local env=$1
-  case "$env" in
-    verily)
+  local env="${1}"
+  case "${env}" in
+    "verily")
       env="prod"
       ;;
-    dev-stable)
+    "dev-stable")
       env="dev"
       ;;
   esac
   echo "https://workbench-app-${env}.verily.com"
 }
+readonly -f get_app_proxy_url
 
 # If the script exits without error let the UI know it completed successfully
 # Otherwise if an error occurred write the line and command that failed to guest attributes.

From 06ae86c03ee8a9cc85e6264701e02dd2d403af80 Mon Sep 17 00:00:00 2001
From: Yu Hu <93544261+yuhuyoyo@users.noreply.github.com>
Date: Mon, 10 Mar 2025 16:02:21 -0400
Subject: [PATCH 03/11] Update startup.sh

---
 startupscript/dataproc/startup.sh | 32 +++++++++++++++++--------------
 1 file changed, 18 insertions(+), 14 deletions(-)

diff --git a/startupscript/dataproc/startup.sh b/startupscript/dataproc/startup.sh
index 3f4386cf..618af140 100644
--- a/startupscript/dataproc/startup.sh
+++ b/startupscript/dataproc/startup.sh
@@ -209,7 +209,7 @@ function get_service_url() {
 readonly -f get_service_url
 
 # Get the app proxy url for the env.
-function get_app_proxy_url() {
+function get_app_proxy_uri() {
   local env="${1}"
   case "${env}" in
     "verily")
@@ -221,7 +221,22 @@ function get_app_proxy_url() {
   esac
   echo "https://workbench-app-${env}.verily.com"
 }
-readonly -f get_app_proxy_url
+readonly -f get_app_proxy_uri
+
+# Get workbench frontend url
+function get_ui_uri() {
+  local env="${1}"
+  if [[ "${env}" == "verily" || "${env}" == "prod" ]]; then
+    echo "https://workbench.verily.com"
+  elif [[ "${env}" == "dev-stable" ]]; then
+    echo "https://workbench-dev.verily.com"
+  else [[ "${env}" == "dev-unstable" || "${env}" == "test" || "${env}" == "staging" ]]
+    echo "https://workbench-${env}.verily.com"
+  else
+    echo "https://workbench.verily.com"
+  fi
+}
+readonly -f get_ui_uri
 
 # If the script exits without error let the UI know it completed successfully
 # Otherwise if an error occurred write the line and command that failed to guest attributes.
@@ -1065,18 +1080,7 @@ fi
   # Configure the original (non-workbench) proxy agent
   ######################################################
     emit "Configuring the original Proxy Agent banner for Google account..."
-    # Map the CLI server to the appropriate UI url
-    if [[ "${TERRA_SERVER}" == *"verily"* ]]; then
-      # Map the CLI server to the appropriate UI url
-      if [[ "${TERRA_SERVER}" == "verily" ]]; then
-        UI_BASE_URL="workbench.verily.com"
-      else
-        UI_BASE_URL="${TERRA_SERVER/verily/terra}-ui-terra.api.verily.com"
-      fi
-    else
-      >&2 echo "ERROR: ${TERRA_SERVER} is not a known verily server."
-      exit 1
-    fi
+    UI_BASE_URL=$(get_ui_url "${TERRA_SERVER}")
     readonly UI_BASE_URL
 
     # The banner.html file contains <style> wrapper tags and a series of CSS styles, and a set of html link elements that we want to modify.

From 1a57ce5db89aea44fde8e5f64b6ae3dd3a8ce9fa Mon Sep 17 00:00:00 2001
From: Yu Hu <93544261+yuhuyoyo@users.noreply.github.com>
Date: Mon, 10 Mar 2025 16:02:58 -0400
Subject: [PATCH 04/11] Update startup.sh

---
 startupscript/dataproc/startup.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/startupscript/dataproc/startup.sh b/startupscript/dataproc/startup.sh
index 618af140..976b5e2c 100644
--- a/startupscript/dataproc/startup.sh
+++ b/startupscript/dataproc/startup.sh
@@ -834,7 +834,7 @@ if [[ "${IS_NON_GOOGLE_ACCOUNT}" == "true" ]]; then
 # Start workbench app proxy agent 
 ###################################
 
-  APP_PROXY="$(get_app_proxy_url "${TERRA_SERVER}")"
+  APP_PROXY="$(get_app_proxy_uri "${TERRA_SERVER}")"
   readonly APP_PROXY
   if [[ -n "${APP_PROXY}" ]]; then
     emit "Using custom Proxy Agent"

From abaeec92b9d87e730939fcd16af74463f2c56c41 Mon Sep 17 00:00:00 2001
From: Yu Hu <93544261+yuhuyoyo@users.noreply.github.com>
Date: Mon, 10 Mar 2025 16:03:36 -0400
Subject: [PATCH 05/11] Update startup.sh

---
 startupscript/dataproc/startup.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/startupscript/dataproc/startup.sh b/startupscript/dataproc/startup.sh
index 976b5e2c..61cfeed3 100644
--- a/startupscript/dataproc/startup.sh
+++ b/startupscript/dataproc/startup.sh
@@ -1080,7 +1080,7 @@ fi
   # Configure the original (non-workbench) proxy agent
   ######################################################
     emit "Configuring the original Proxy Agent banner for Google account..."
-    UI_BASE_URL=$(get_ui_url "${TERRA_SERVER}")
+    UI_BASE_URL=$(get_ui_uri "${TERRA_SERVER}")
     readonly UI_BASE_URL
 
     # The banner.html file contains <style> wrapper tags and a series of CSS styles, and a set of html link elements that we want to modify.

From 10a7a2ad666f7b35b72db17b5d412efdaf0f742a Mon Sep 17 00:00:00 2001
From: Yu Hu <93544261+yuhuyoyo@users.noreply.github.com>
Date: Fri, 14 Mar 2025 08:07:10 -0400
Subject: [PATCH 06/11] Update startup.sh

---
 startupscript/dataproc/startup.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/startupscript/dataproc/startup.sh b/startupscript/dataproc/startup.sh
index 61cfeed3..26b3d07d 100644
--- a/startupscript/dataproc/startup.sh
+++ b/startupscript/dataproc/startup.sh
@@ -219,7 +219,7 @@ function get_app_proxy_uri() {
       env="dev"
       ;;
   esac
-  echo "https://workbench-app-${env}.verily.com"
+  echo "https://workbench-app-${env}.verily.com/api/proxy/"
 }
 readonly -f get_app_proxy_uri
 

From bc52e70694fbad9c7167253088ec6ae129dadc77 Mon Sep 17 00:00:00 2001
From: Yu Hu <93544261+yuhuyoyo@users.noreply.github.com>
Date: Fri, 14 Mar 2025 08:09:29 -0400
Subject: [PATCH 07/11] Update startup.sh

---
 startupscript/dataproc/startup.sh | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/startupscript/dataproc/startup.sh b/startupscript/dataproc/startup.sh
index 26b3d07d..09d00593 100644
--- a/startupscript/dataproc/startup.sh
+++ b/startupscript/dataproc/startup.sh
@@ -219,7 +219,7 @@ function get_app_proxy_uri() {
       env="dev"
       ;;
   esac
-  echo "https://workbench-app-${env}.verily.com/api/proxy/"
+  echo "https://workbench-app-${env}.verily.com/"
 }
 readonly -f get_app_proxy_uri
 
@@ -839,7 +839,7 @@ if [[ "${IS_NON_GOOGLE_ACCOUNT}" == "true" ]]; then
   if [[ -n "${APP_PROXY}" ]]; then
     emit "Using custom Proxy Agent"
     RESOURCE_ID="$(get_metadata_value "instance/attributes/terra-resource-id")"
-    NEW_PROXY="${APP_PROXY}"
+    NEW_PROXY="${APP_PROXY}/api/proxy"
     NEW_PROXY_URL="${RESOURCE_ID}.${APP_PROXY}"
     readonly RESOURCE_ID
     readonly NEW_PROXY

From 6bdb0edcccf7ad33b11e9c88a264060a94c05fb4 Mon Sep 17 00:00:00 2001
From: Yu Hu <yuhuyoyo@verily.health>
Date: Tue, 25 Mar 2025 14:31:54 -0400
Subject: [PATCH 08/11] detect active ssh session

---
 startupscript/butane/idle-shutdown.sh     | 8 +++++++-
 startupscript/butane/probe-user-access.sh | 6 ++++++
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/startupscript/butane/idle-shutdown.sh b/startupscript/butane/idle-shutdown.sh
index 6ec52eb7..3964453e 100644
--- a/startupscript/butane/idle-shutdown.sh
+++ b/startupscript/butane/idle-shutdown.sh
@@ -38,11 +38,17 @@ readonly LAST_BOOT_TIME
 # Get the last time the VM was active. Default to 0 if not set.
 LAST_ACTIVE="$(get_guest_attribute "last-active/cpu" "0")"
 LAST_ACTIVE_PROXY="$(get_guest_attribute "last-active/proxy" "0")"
+LAST_ACTIVE_SSH="$(get_guest_attribute "last-active/ssh" "0")"
 
-# get the latest time between the two last active timestamps
+# get the latest active timestamps
 if (( LAST_ACTIVE < LAST_ACTIVE_PROXY )); then
     LAST_ACTIVE="${LAST_ACTIVE_PROXY}"
 fi
+
+if (( LAST_ACTIVE < LAST_ACTIVE_SSH )); then
+    LAST_ACTIVE="${LAST_ACTIVE_SSH}"
+fi
+
 # get the latest time between the last boot time and the last active time. The last active timestamp could still be really old when the VM is rebooted and
 # the last active timestamp is not updated yet.
 if (( LAST_ACTIVE < LAST_BOOT_TIME )); then
diff --git a/startupscript/butane/probe-user-access.sh b/startupscript/butane/probe-user-access.sh
index 2467f44d..b05547e5 100644
--- a/startupscript/butane/probe-user-access.sh
+++ b/startupscript/butane/probe-user-access.sh
@@ -39,3 +39,9 @@ else
     NOW="$(date +'%s')"
     set_metadata "last-active/cpu" "${NOW}"
 fi
+
+if ps auxwww | grep sshd:; then
+  echo "Detect an active ssh session"
+  NOW="$(date +'%s')"
+  set_metadata "last-active/ssh" "${NOW}"
+fi

From 1709ee9354b06a5be446c0ba654a67ea19e77fba Mon Sep 17 00:00:00 2001
From: Yu Hu <yuhuyoyo@verily.health>
Date: Tue, 25 Mar 2025 14:33:18 -0400
Subject: [PATCH 09/11] restore

---
 startupscript/dataproc/startup.sh | 131 +++++++++++++++++++++---------
 1 file changed, 92 insertions(+), 39 deletions(-)

diff --git a/startupscript/dataproc/startup.sh b/startupscript/dataproc/startup.sh
index 09d00593..023b9994 100644
--- a/startupscript/dataproc/startup.sh
+++ b/startupscript/dataproc/startup.sh
@@ -197,13 +197,13 @@ function set_guest_attributes() {
 # Map the CLI server to appropriate service path
 function get_service_url() {
   case "$1" in
-    "verily") echo "https://workbench.verily.com/api/$2" ;;
-    "dev-stable") echo "https://workbench-dev.verily.com/api/$2" ;;
-    "dev-unstable") echo "https://workbench-dev-unstable.verily.com/api/$2" ;;
-    "test") echo "https://workbench-test.verily.com/api/$2" ;;
-    "staging") echo "https://workbench-staging.verily.com/api/$2" ;;
-    "prod") echo "https://workbench.verily.com/api/$2" ;;
-    *) return 1 ;;
+  "verily") echo "https://workbench.verily.com/api/$2" ;;
+  "dev-stable") echo "https://workbench-dev.verily.com/api/$2" ;;
+  "dev-unstable") echo "https://workbench-dev-unstable.verily.com/api/$2" ;;
+  "test") echo "https://workbench-test.verily.com/api/$2" ;;
+  "staging") echo "https://workbench-staging.verily.com/api/$2" ;;
+  "prod") echo "https://workbench.verily.com/api/$2" ;;
+  *) return 1 ;;
   esac
 }
 readonly -f get_service_url
@@ -212,14 +212,14 @@ readonly -f get_service_url
 function get_app_proxy_uri() {
   local env="${1}"
   case "${env}" in
-    "verily")
-      env="prod"
-      ;;
-    "dev-stable")
-      env="dev"
-      ;;
+  "verily")
+    env="prod"
+    ;;
+  "dev-stable")
+    env="dev"
+    ;;
   esac
-  echo "https://workbench-app-${env}.verily.com/"
+  echo "workbench-app-${env}.verily.com"
 }
 readonly -f get_app_proxy_uri
 
@@ -230,7 +230,7 @@ function get_ui_uri() {
     echo "https://workbench.verily.com"
   elif [[ "${env}" == "dev-stable" ]]; then
     echo "https://workbench-dev.verily.com"
-  else [[ "${env}" == "dev-unstable" || "${env}" == "test" || "${env}" == "staging" ]]
+  elif [[ "${env}" == "dev-unstable" || "${env}" == "test" || "${env}" == "staging" ]]; then
     echo "https://workbench-${env}.verily.com"
   else
     echo "https://workbench.verily.com"
@@ -293,7 +293,7 @@ function install_nextflow() {
 readonly -f install_nextflow
 
 #######################################
-### Begin environment setup 
+### Begin environment setup
 #######################################
 
 # Let the UI know the script has started
@@ -627,7 +627,7 @@ EOF
 
 emit "Setting up git integration..."
 
-# Create the user SSH directory 
+# Create the user SSH directory
 ${RUN_AS_LOGIN_USER} "mkdir -p ${USER_SSH_DIR} --mode 0700"
 
 # Get the user's SSH key from Workbench, and if set, write it to the user's .ssh directory
@@ -824,29 +824,79 @@ chown "${LOGIN_USER}:${LOGIN_USER}" "${USER_BASH_PROFILE}"
 
 # TODO(BENCH-2612): use workbench CLI instead to get user profile.
 IS_NON_GOOGLE_ACCOUNT="$(curl "${USER_SERVICE_URL}/api/profile?path=non_google_account" \
-                    -H "accept: application/json" -H "Authorization: Bearer $(wb auth print-access-token)" \
+                    -H "accept: application/json" -H "Authorization: Bearer $(gcloud auth print-access-token)" \
                   | jq '.value')"
 readonly IS_NON_GOOGLE_ACCOUNT
 
-APP_PROXY=""
-if [[ "${IS_NON_GOOGLE_ACCOUNT}" == "true" ]]; then
+ENABLE_VWB_PROXY="$(get_metadata_value "instance/attributes/enable-dataproc-vwb-proxy")"
+readonly ENABLE_VWB_PROXY
+# Retrieve Workbench proxy agent image
+PROXY_AGENT_IMAGE="$(get_metadata_value "instance/attributes/proxy-agent-image")"
+readonly PROXY_AGENT_IMAGE
+APP_PROXY="$(get_app_proxy_uri "${TERRA_SERVER}")"
+readonly APP_PROXY
+RESOURCE_ID="$(get_metadata_value "instance/attributes/terra-resource-id")"
+readonly RESOURCE_ID
+# Retrieve app proxy user facing URL
+APP_PROXY_URL="https://${RESOURCE_ID}.${APP_PROXY}"
+readonly APP_PROXY_URL
+# Retrieve app proxy service facing URL
+PROXY_SERVICE_URL="$(get_service_url "${TERRA_SERVER}" "proxy")/"
+readonly PROXY_SERVICE_URL
+
+# Define constants for proxy types
+PROXY_TYPE_VWB="VWB_PROXY"
+PROXY_TYPE_GOOGLE="GOOGLE_PROXY"
+PROXY_TYPE_GOOGLE_AGENT_WITH_VWB_PROXY="GOOGLE_AGENT_WITH_VWB_PROXY"
+
+# Decide proxy type
+PROXY_TYPE=""
+if [[ ${ENABLE_VWB_PROXY} == "true" && -n "${PROXY_AGENT_IMAGE}" ]]; then
+  PROXY_TYPE="${PROXY_TYPE_VWB}"
+elif [[ ${IS_NON_GOOGLE_ACCOUNT} == "true" ]]; then
+  PROXY_TYPE="${PROXY_TYPE_GOOGLE_AGENT_WITH_VWB_PROXY}"
+else
+  PROXY_TYPE="${PROXY_TYPE_GOOGLE}"
+fi
+
+if [[ "${PROXY_TYPE}" != "${PROXY_TYPE_GOOGLE}" ]]; then
 ###################################
 # Start workbench app proxy agent 
 ###################################
+    if [[ "${PROXY_TYPE}" == "${PROXY_TYPE_VWB}" ]]; then
+      emit "Using Workbench Proxy"
+      ENABLE_MONITORING_SCRIPT="$(get_metadata_value "instance/attributes/enable-dataproc-monitoring-script")"
+      readonly ENABLE_MONITORING_SCRIPT
+
+      cat <<EOF >"${WORKBENCH_PROXY_AGENT_SERVICE}"
+[Unit]
+Description=Workbench Docker Proxy Agent
+After=docker.service
+Requires=docker.service
+
+[Service]
+ExecStartPre=/bin/bash -c 'while [ ! -f ${PROXY_AGENT_BANNER} ]; do sleep 5; done'
+ExecStart=/bin/bash -c '/usr/bin/docker run \\
+      --detach \\
+      --name "proxy-agent" \\
+      --restart=unless-stopped \\
+      --net=host "${PROXY_AGENT_IMAGE}" \\
+      --proxy="${PROXY_SERVICE_URL}" \\
+      --host="localhost:8123" \\
+      --compute-platform="GCP" \\
+      --shim-path="websocket-shim" \\
+      --rewrite-websocket-host="true" \\
+      --backend="${RESOURCE_ID}" \\
+      --session-cookie-name="_xsrf" \\
+      --inject-banner="\$(cat ${PROXY_AGENT_BANNER})" \\
+      --enable-monitoring-script="${ENABLE_MONITORING_SCRIPT:-false}"'
 
-  APP_PROXY="$(get_app_proxy_uri "${TERRA_SERVER}")"
-  readonly APP_PROXY
-  if [[ -n "${APP_PROXY}" ]]; then
-    emit "Using custom Proxy Agent"
-    RESOURCE_ID="$(get_metadata_value "instance/attributes/terra-resource-id")"
-    NEW_PROXY="${APP_PROXY}/api/proxy"
-    NEW_PROXY_URL="${RESOURCE_ID}.${APP_PROXY}"
-    readonly RESOURCE_ID
-    readonly NEW_PROXY
-    readonly NEW_PROXY_URL
-
-    # Create a systemd service to start the workbench app proxy.
-    cat << EOF > "${WORKBENCH_PROXY_AGENT_SERVICE}"
+[Install]
+WantedBy=multi-user.target
+EOF
+    else
+      emit "Using Google Proxy Agent with Workbench Proxy"
+      cat <<EOF >"${WORKBENCH_PROXY_AGENT_SERVICE}"
 [Unit]
 Description=Workbench App Proxy Agent Service
 
@@ -857,13 +907,16 @@ Restart=on-failure
 [Install]
 WantedBy=multi-user.target
 EOF
+    fi
 
     # Enable and start the workbench app proxy agent service
     systemctl daemon-reload
     systemctl enable "${WORKBENCH_PROXY_AGENT_SERVICE_NAME}"
-    systemctl start "${WORKBENCH_PROXY_AGENT_SERVICE_NAME}"
+    # Start the service in non-blocking mode to allow startup script to continue without waiting
+    systemctl --no-block start "${WORKBENCH_PROXY_AGENT_SERVICE_NAME}"
     emit "Workbench Proxy Agent service started"
-  fi
+else
+  emit "Using Google Proxy"
 fi
 
 ###################################
@@ -1074,7 +1127,7 @@ fi
     emit "Starting Hail install script..."
     ${RUN_PYTHON} ${HAIL_SCRIPT_PATH}
   fi
-  
+
   if [[ "${IS_NON_GOOGLE_ACCOUNT}" == "false" ]]; then
   ######################################################
   # Configure the original (non-workbench) proxy agent
@@ -1128,11 +1181,11 @@ EOF
   sed -i -e "/c.GCSContentsManager/d" -e "/CombinedContentsManager/d" "${JUPYTER_CONFIG}"
   echo "c.FileContentsManager.root_dir = '${USER_HOME_DIR}'" >> "${JUPYTER_CONFIG}"
 
-  if [[ -n "${APP_PROXY}" ]] && [[ "${IS_NON_GOOGLE_ACCOUNT}" == "true" ]]; then
+  if [[ "${PROXY_TYPE}" != "${PROXY_TYPE_GOOGLE}" ]]; then
     cat << EOF >> "${JUPYTER_CONFIG}"
-c.NotebookApp.allow_origin_pat += "|(https?://)?(https://${NEW_PROXY_URL})"
+c.NotebookApp.allow_origin_pat += "|(https?://)?(${APP_PROXY_URL})"
 c.NotebookApp.allow_remote_access= True
-c.NotebookApp.local_hostnames.append('${APP_PROXY}')
+c.NotebookApp.local_hostnames.append('${PROXY_SERVICE_URL}')
 EOF
   fi
   # Add a marker for the Workbench-specific customizations

From c8072d851fdbd9a9ea1073bc9c74bdbe2d1bdbb7 Mon Sep 17 00:00:00 2001
From: Yu Hu <yuhuyoyo@verily.health>
Date: Tue, 25 Mar 2025 14:49:27 -0400
Subject: [PATCH 10/11] grep active sshd session

---
 startupscript/butane/probe-user-access.sh | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/startupscript/butane/probe-user-access.sh b/startupscript/butane/probe-user-access.sh
index b05547e5..acadd182 100644
--- a/startupscript/butane/probe-user-access.sh
+++ b/startupscript/butane/probe-user-access.sh
@@ -40,8 +40,9 @@ else
     set_metadata "last-active/cpu" "${NOW}"
 fi
 
-if ps auxwww | grep sshd:; then
+if ps auxwww | grep sshd-session | grep -v grep | grep '@'; then
   echo "Detect an active ssh session"
   NOW="$(date +'%s')"
   set_metadata "last-active/ssh" "${NOW}"
 fi
+

From 358debc2dd6da351368da6a07d6642c80828165c Mon Sep 17 00:00:00 2001
From: Yu Hu <yuhuyoyo@verily.health>
Date: Tue, 25 Mar 2025 14:53:36 -0400
Subject: [PATCH 11/11] restore

---
 startupscript/dataproc/startup.sh | 1 +
 1 file changed, 1 insertion(+)
 mode change 100644 => 100755 startupscript/dataproc/startup.sh

diff --git a/startupscript/dataproc/startup.sh b/startupscript/dataproc/startup.sh
old mode 100644
new mode 100755
index 023b9994..1e435015
--- a/startupscript/dataproc/startup.sh
+++ b/startupscript/dataproc/startup.sh
@@ -1329,3 +1329,4 @@ fi
 emit "SUCCESS: Gitignore installed at ${INSTALLED_GITIGNORE}"
 
 # TODO: Pending CLI support in PF-2865 for setting cluster metadata for testing once we have CLI support.
+