beaker-session.zsh

beaker-session() {
  # Defaults
  local cluster="ai2/hammond"

  # Parse flags
  while (( $# > 0 )); do
    case "$1" in
      -h|--help)
        echo "beaker-session [-h|--help] [-c|--cluster CLUSTER|--cluster=CLUSTER]"
        return 0
        ;;
      -c|--cluster)
        [[ $# -ge 2 ]] || { echo "Error: --cluster requires an argument" >&2; return 1; }
        cluster="$2"
        shift 2
        ;;
      --cluster=*)
        cluster="${1#*=}"
        shift
        ;;
      --)
        shift
        break
        ;;
      -*)
        echo "Error: Unknown option: $1" >&2
        return 1
        ;;
      *)
        echo "Error: beaker-session takes no positional arguments" >&2
        return 1
        ;;
    esac
  done

  (( $# == 0 )) || { echo "Error: beaker-session takes no positional arguments" >&2; return 1; }

  local author session node

  author="$(beaker account whoami --format=json | jq -r '.[].name')" || return $?
  [[ -n "$author" ]] || { echo "Error: Could not determine beaker author." >&2; return 1; }

  # Look for an existing session
  session="$(
    beaker session list --author="$author" --cluster="$cluster" --format=json \
      | jq -r '.[].id' | head -n 1
  )" || return $?

  if [[ -n "$session" ]]; then
    echo "Using existing session: $session"
  else
    # Create a new detached session
    beaker session create --detach --remote --bare \
      --cluster "$cluster" \
      --workspace "ai2/open-instruct-dev" \
      --budget "ai2/allennlp"

    session="$(
      beaker session list --author="$author" --cluster="$cluster" --format=json \
        | jq -r '.[].id' | head -n 1
    )" || return $?

    [[ -n "$session" ]] || { echo "Error: Failed to create a new session." >&2; return 1; }
    echo "Created new session: $session"
  fi

  node="$(
    beaker session get "$session" --format=json \
      | jq -r '.[].session.envVars[] | select(.name=="BEAKER_NODE_HOSTNAME") | .value'
  )" || return $?

  [[ -n "$node" ]] || { echo "Error: Could not determine node hostname for session $session" >&2; return 1\
; }

  # Always use mosh (no fallback)
  mosh "$author@$node" -- beaker session attach "$session"
}