ai-code-app/etc/commands/server-command/restart-work-server.sh

#!/bin/sh

set -eu

SCRIPT_DIR=$(CDPATH= cd -- "$(dirname -- "$0")" && pwd)
REPO_ROOT=$(CDPATH= cd -- "$SCRIPT_DIR/../../.." && pwd)
COMPOSE_FILE="$REPO_ROOT/etc/servers/work-server/docker-compose.yml"
PROXY_SERVICE="${WORK_SERVER_PROXY_SERVICE:-work-server}"
PROXY_CONTAINER="${WORK_SERVER_PROXY_CONTAINER:-work-server}"
BLUE_SERVICE="${WORK_SERVER_BLUE_SERVICE:-work-server-blue}"
GREEN_SERVICE="${WORK_SERVER_GREEN_SERVICE:-work-server-green}"
BLUE_CONTAINER="${WORK_SERVER_BLUE_CONTAINER:-work-server-blue}"
GREEN_CONTAINER="${WORK_SERVER_GREEN_CONTAINER:-work-server-green}"
ACTIVE_SLOT_FILE="${WORK_SERVER_ACTIVE_SLOT_FILE:-$REPO_ROOT/etc/servers/work-server/.docker/runtime/active-slot}"
PROXY_CONFIG_FILE="${WORK_SERVER_PROXY_CONFIG_FILE:-$REPO_ROOT/etc/servers/work-server/.docker/proxy/default.conf}"
HEALTH_ENDPOINT="${WORK_SERVER_HEALTH_ENDPOINT:-http://127.0.0.1:3100/health}"
RUNTIME_ENDPOINT="${WORK_SERVER_RUNTIME_ENDPOINT:-http://127.0.0.1:3100/api/runtime}"
PREVIOUS_SLOT_DRAIN_TIMEOUT_SECONDS="${WORK_SERVER_PREVIOUS_SLOT_DRAIN_TIMEOUT_SECONDS:-900}"

cd "$REPO_ROOT"

mkdir -p "$(dirname "$ACTIVE_SLOT_FILE")" "$(dirname "$PROXY_CONFIG_FILE")"

read_active_slot() {
  if [ -f "$ACTIVE_SLOT_FILE" ]; then
    SLOT=$(tr -d '[:space:]' <"$ACTIVE_SLOT_FILE")
    if [ "$SLOT" = "blue" ] || [ "$SLOT" = "green" ]; then
      printf '%s' "$SLOT"
      return 0
    fi
  fi

  printf 'blue'
}

container_is_running() {
  CONTAINER_NAME="$1"
  STATUS=$(docker inspect -f '{{.State.Status}}' "$CONTAINER_NAME" 2>/dev/null || true)
  [ "$STATUS" = "running" ]
}

resolve_active_slot() {
  SLOT=$(read_active_slot)

  if [ "$SLOT" = "blue" ] && ! container_is_running "$BLUE_CONTAINER" && container_is_running "$GREEN_CONTAINER"; then
    printf 'green'
    return 0
  fi

  if [ "$SLOT" = "green" ] && ! container_is_running "$GREEN_CONTAINER" && container_is_running "$BLUE_CONTAINER"; then
    printf 'blue'
    return 0
  fi

  printf '%s' "$SLOT"
}

write_proxy_config() {
  SLOT="$1"
  TARGET_CONTAINER="$BLUE_CONTAINER"

  if [ "$SLOT" = "green" ]; then
    TARGET_CONTAINER="$GREEN_CONTAINER"
  fi

  cat >"$PROXY_CONFIG_FILE" <<EOF2
server {
  listen 3100;
  server_name _;

  location /ws/chat {
    proxy_http_version 1.1;
    proxy_set_header Host \$host;
    proxy_set_header X-Forwarded-For \$proxy_add_x_forwarded_for;
    proxy_set_header X-Forwarded-Proto \$scheme;
    proxy_set_header X-Forwarded-Host \$host;
    proxy_set_header X-Forwarded-Port \$server_port;
    proxy_set_header Upgrade \$http_upgrade;
    proxy_set_header Connection "upgrade";
    proxy_pass http://$TARGET_CONTAINER:3100;
  }

  location / {
    proxy_http_version 1.1;
    proxy_set_header Host \$host;
    proxy_set_header X-Forwarded-For \$proxy_add_x_forwarded_for;
    proxy_set_header X-Forwarded-Proto \$scheme;
    proxy_set_header X-Forwarded-Host \$host;
    proxy_set_header X-Forwarded-Port \$server_port;
    proxy_set_header Connection "";
    proxy_pass http://$TARGET_CONTAINER:3100;
  }
}
EOF2
}

wait_for_container_health() {
  TARGET_CONTAINER="$1"
  ATTEMPT=0

  while [ "$ATTEMPT" -lt 60 ]; do
    if docker exec "$TARGET_CONTAINER" node -e "fetch(process.argv[1]).then(async (response) => { if (!response.ok) process.exit(1); process.stdout.write(await response.text()); }).catch(() => process.exit(1));" "$HEALTH_ENDPOINT" >/dev/null 2>&1; then
      return 0
    fi

    ATTEMPT=$((ATTEMPT + 1))
    sleep 2
  done

  echo "health check failed for $TARGET_CONTAINER" >&2
  return 1
}

read_runtime_value() {
  TARGET_CONTAINER="$1"
  FIELD_NAME="$2"

  docker exec "$TARGET_CONTAINER" node -e "fetch(process.argv[1]).then((response) => response.json()).then((payload) => { const value = payload?.[process.argv[2]]; if (typeof value === 'boolean') { process.stdout.write(value ? 'true' : 'false'); return; } if (value == null) { process.stdout.write(''); return; } process.stdout.write(String(value)); }).catch(() => process.exit(1));" "$RUNTIME_ENDPOINT" "$FIELD_NAME"
}

set_container_draining() {
  TARGET_CONTAINER="$1"
  DRAINING_VALUE="$2"

  docker exec "$TARGET_CONTAINER" node -e "fetch(process.argv[1], { method: 'POST', headers: { 'content-type': 'application/json' }, body: JSON.stringify({ draining: process.argv[2] === 'true' }) }).then((response) => { if (!response.ok) process.exit(1); }).catch(() => process.exit(1));" "${RUNTIME_ENDPOINT}/drain" "$DRAINING_VALUE"
}

wait_for_previous_slot_drain() {
  TARGET_CONTAINER="$1"
  ELAPSED=0

  while [ "$ELAPSED" -lt "$PREVIOUS_SLOT_DRAIN_TIMEOUT_SECONDS" ]; do
    ACTIVE_COUNT=$(read_runtime_value "$TARGET_CONTAINER" activeChatRequestCount 2>/dev/null || printf '0')
    QUEUED_COUNT=$(read_runtime_value "$TARGET_CONTAINER" queuedChatRequestCount 2>/dev/null || printf '0')

    if [ "${ACTIVE_COUNT:-0}" = "0" ] && [ "${QUEUED_COUNT:-0}" = "0" ]; then
      return 0
    fi

    sleep 2
    ELAPSED=$((ELAPSED + 2))
  done

  echo "drain timeout reached for $TARGET_CONTAINER" >&2
  return 1
}

ensure_proxy_running() {
  docker compose -f "$COMPOSE_FILE" up -d --no-deps "$PROXY_SERVICE" >/dev/null
  docker exec "$PROXY_CONTAINER" nginx -s reload >/dev/null
}

ACTIVE_SLOT=$(resolve_active_slot)
TARGET_SLOT="green"
TARGET_SERVICE="$GREEN_SERVICE"
TARGET_CONTAINER="$GREEN_CONTAINER"
PREVIOUS_SERVICE="$BLUE_SERVICE"
PREVIOUS_CONTAINER="$BLUE_CONTAINER"

if [ "$ACTIVE_SLOT" = "green" ]; then
  TARGET_SLOT="blue"
  TARGET_SERVICE="$BLUE_SERVICE"
  TARGET_CONTAINER="$BLUE_CONTAINER"
  PREVIOUS_SERVICE="$GREEN_SERVICE"
  PREVIOUS_CONTAINER="$GREEN_CONTAINER"
fi

docker compose -f "$COMPOSE_FILE" up -d --build --force-recreate --no-deps "$TARGET_SERVICE"
wait_for_container_health "$TARGET_CONTAINER"
write_proxy_config "$TARGET_SLOT"
ensure_proxy_running
printf '%s\n' "$TARGET_SLOT" >"$ACTIVE_SLOT_FILE"

if [ "$PREVIOUS_SERVICE" != "$TARGET_SERVICE" ]; then
  set_container_draining "$PREVIOUS_CONTAINER" true
  wait_for_previous_slot_drain "$PREVIOUS_CONTAINER"
  docker compose -f "$COMPOSE_FILE" up -d --build --force-recreate --no-deps "$PREVIOUS_SERVICE"
  wait_for_container_health "$PREVIOUS_CONTAINER"
fi

printf 'work-server zero-downtime switch completed: %s -> %s\n' "$ACTIVE_SLOT" "$TARGET_SLOT"