Waiting for a sandbox

#!/usr/bin/env bash
# Get a sandbox, waiting in the queue when Atlas Labs is full.
# POST /sandboxes answers 201 with your keys, or 202 with a queue ticket when every slot is taken.
# Poll the ticket until a poll answers 201: that body is the only time the keys appear, and a
# ready ticket must be claimed within 120 s or it is forfeited (410 ticket_forfeited).
# Needs: curl, jq.
set -euo pipefail
BASE="${ATLAS_BASE:-https://automationatlas.io/labs/api/v1}"
ORIGIN=$(printf '%s' "$BASE" | sed -E 's#^(https?://[^/]+).*#\1#') # `poll` is a path, not a URL
headers=$(mktemp) body=$(mktemp)
trap 'rm -f "$headers" "$body"' EXIT

# Run one request; sets $status, leaves the response in $body and its headers in $headers.
call() { status=$(curl -sS -D "$headers" -o "$body" -w '%{http_code}' "$@"); }
retry_after() { awk -F': *' 'tolower($1) == "retry-after" { gsub(/\r/, "", $2); print $2 }' "$headers"; }

for attempt in 1 2 3 4 5; do
  call -X POST "$BASE/sandboxes" -H "Content-Type: application/json" -d '{"preset":"starter"}'
  # 503 queue_full: the queue itself is full. Retry-After says when to try again (30 s).
  if [ "$status" = 503 ] && [ "$attempt" -lt 5 ]; then
    pause=$(retry_after); sleep "${pause:-30}"; continue
  fi
  break
done

if [ "$status" = 202 ]; then
  poll=$(jq -r .data.poll "$body")
  echo "queued at position $(jq -r .data.position "$body"), about $(jq -r .data.eta_s "$body") seconds"
  while :; do
    pause=$(retry_after); pause=${pause:-3}
    [ "$pause" -ge 3 ] 2>/dev/null || pause=3 # unauthenticated calls share 30 per minute per IP group
    sleep "$pause"
    call "$ORIGIN$poll"
    case "$status" in
      202) echo "position $(jq -r .data.position "$body")" ;;
      429 | 5??) ;; # keep the place, wait Retry-After
      *) break ;;    # 201 keys, 410 forfeited, 404 unknown or already claimed
    esac
  done
fi

if [ "$status" != 201 ]; then
  echo "HTTP $status: $(cat "$body")" >&2
  exit 1
fi

# Owner-only file. Do not commit it.
(umask 077
 printf 'export ATLAS_FACTORY_KEY=%s\nexport ATLAS_OFFICE_KEY=%s\n' \
   "$(jq -r .data.factory_key "$body")" "$(jq -r .data.office_key "$body")" > atlas-keys.env)
chmod 600 atlas-keys.env  # umask only applies to a new file
jq '{sandbox_id: .data.sandbox_id, name: .data.public_name, expires_at: .data.expires_at}' "$body"
echo "Keys saved to atlas-keys.env. Run: source atlas-keys.env"

// Get a sandbox, waiting in the queue when Atlas Labs is full (Node 20+, no dependencies).
// POST /sandboxes answers 201 with your keys, or 202 with a queue ticket when every slot is taken.
// Poll the ticket until a poll answers 201: that body is the only time the keys appear, and a
// ready ticket must be claimed within 120 s or it is forfeited (410 ticket_forfeited).
import { chmodSync, writeFileSync } from "node:fs";

const BASE = process.env.ATLAS_BASE ?? "https://automationatlas.io/labs/api/v1";
const ORIGIN = new URL(BASE).origin; // `poll` is a path such as /labs/api/v1/queue/tkt_..., not a URL
const POLL_FLOOR_MS = 3000; // unauthenticated calls share 30 per minute per IP group
const sleep = (ms) => new Promise((resolve) => setTimeout(resolve, ms));
const retryAfterMs = (res, fallbackMs) => {
  const seconds = Number(res.headers.get("retry-after"));
  return Number.isFinite(seconds) && seconds > 0 ? seconds * 1000 : fallbackMs;
};

async function create(body) {
  for (let attempt = 1; ; attempt += 1) {
    const res = await fetch(`${BASE}/sandboxes`, {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify(body),
    });
    // 503 queue_full: the queue itself is full. Retry-After says when to try again (30 s).
    if (res.status === 503 && attempt < 5) {
      await sleep(retryAfterMs(res, 30_000));
      continue;
    }
    return res;
  }
}

let res = await create({ preset: "starter" });
let json = await res.json();

if (res.status === 202) {
  const { poll } = json.data;
  console.log("queued at position", json.data.position, "about", json.data.eta_s, "seconds");
  for (;;) {
    await sleep(Math.max(POLL_FLOOR_MS, retryAfterMs(res, 0)));
    res = await fetch(ORIGIN + poll);
    json = await res.json().catch(() => ({}));
    if (res.status === 202) {
      console.log("position", json.data.position);
      continue;
    }
    if (res.status === 429 || res.status >= 500) continue; // keep the place, wait Retry-After
    break; // 201 keys, 410 forfeited, 404 unknown or already claimed
  }
}
if (res.status !== 201) {
  const e = json.error ?? {};
  throw new Error(`HTTP ${res.status} ${e.code}: ${e.hint ?? e.message}`);
}
const sandbox = json.data;

// Owner-only file. Do not commit it.
writeFileSync(
  "atlas-keys.env",
  `export ATLAS_FACTORY_KEY=${sandbox.factory_key}\nexport ATLAS_OFFICE_KEY=${sandbox.office_key}\n`,
  { mode: 0o600 },
);
// `mode` only applies when the file is created; tighten one that already existed.
chmodSync("atlas-keys.env", 0o600);
console.log(sandbox.sandbox_id, sandbox.public_name, "expires", sandbox.expires_at);
console.log("Keys saved to atlas-keys.env. Run: source atlas-keys.env");

"""Get a sandbox, waiting in the queue when Atlas Labs is full.

POST /sandboxes answers 201 with your keys, or 202 with a queue ticket when every slot is taken.
Poll the ticket until a poll answers 201: that body is the only time the keys appear, and a
ready ticket must be claimed within 120 s or it is forfeited (410 ticket_forfeited).

Needs: pip install requests
"""
import os
import time
from urllib.parse import urlsplit

import requests

BASE = os.environ.get("ATLAS_BASE", "https://automationatlas.io/labs/api/v1")
parts = urlsplit(BASE)
ORIGIN = f"{parts.scheme}://{parts.netloc}"  # `poll` is a path such as /labs/api/v1/queue/tkt_..., not a URL
POLL_FLOOR_S = 3  # unauthenticated calls share 30 per minute per IP group


def retry_after(res, fallback):
    try:
        seconds = float(res.headers.get("Retry-After", ""))
        return seconds if seconds > 0 else fallback
    except ValueError:
        return fallback


def create(body):
    for attempt in range(1, 6):
        res = requests.post(f"{BASE}/sandboxes", json=body, timeout=30)
        # 503 queue_full: the queue itself is full. Retry-After says when to try again (30 s).
        if res.status_code == 503 and attempt < 5:
            time.sleep(retry_after(res, 30))
            continue
        return res
    return res


res = create({"preset": "starter"})

if res.status_code == 202:
    data = res.json()["data"]
    poll = data["poll"]
    print("queued at position", data["position"], "about", data["eta_s"], "seconds")
    while True:
        time.sleep(max(POLL_FLOOR_S, retry_after(res, 0)))
        res = requests.get(ORIGIN + poll, timeout=30)
        if res.status_code == 202:
            print("position", res.json()["data"]["position"])
            continue
        if res.status_code == 429 or res.status_code >= 500:
            continue  # keep the place, wait Retry-After
        break  # 201 keys, 410 forfeited, 404 unknown or already claimed

if res.status_code != 201:
    err = res.json().get("error", {})
    raise SystemExit(f"HTTP {res.status_code} {err.get('code')}: {err.get('hint') or err.get('message')}")
sandbox = res.json()["data"]

# Owner-only file. Do not commit it.
fd = os.open("atlas-keys.env", os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
with os.fdopen(fd, "w") as f:
    f.write(f"export ATLAS_FACTORY_KEY={sandbox['factory_key']}\n")
    f.write(f"export ATLAS_OFFICE_KEY={sandbox['office_key']}\n")
os.chmod("atlas-keys.env", 0o600)  # the mode above only applies to a new file

print(sandbox["sandbox_id"], sandbox["public_name"], "expires", sandbox["expires_at"])
print("Keys saved to atlas-keys.env. Run: source atlas-keys.env")