Leveraging multiple, repository-specific OpenAI Codex API Keys with Visual Studio Code on macOS

Be sure to update the following based on your API Keys:

OPENAI_GENERAL_ITEM_NAME="${OPENAI_GENERAL_ITEM_NAME:-Example OpenAI General Key}"

# Update `_key_names` to match your masked API key suffixes.
local -A _key_names=(
    [kPqX]="Microsoft 365 Reset"
    [vL7m]="Mac Health Check"
    [rT9p]="DDM OS Reminder"
    [xY2k]="Setup Your Mac"
    [bF8v]="General"
)

# These functions will need to be manually added to your `~/.zshrc` file to work properly.
# See: [bz please](https://snelson.us/2022/05/bz-please/)

# --- OpenAI Codex: Start -----------------------------------------------------

OPENAI_GENERAL_ITEM_NAME="${OPENAI_GENERAL_ITEM_NAME:-Example OpenAI General Key}"

_oaki () {   # [o]pen[a]i [k]ey [i]dentity — resolve Name from last-4 of key
    local key="${1:-${OPENAI_API_KEY:-}}"
    [[ -n "$key" ]] || {
        printf "_oaki: no key provided and OPENAI_API_KEY is unset\n" >&2
        return 1
    }

    local -A _key_names=(
        [kPqX]="Microsoft 365 Reset"
        [vL7m]="Mac Health Check"
        [rT9p]="DDM OS Reminder"
        [xY2k]="Setup Your Mac"
        [bF8v]="General"
    )

    local suffix="${key[-4,-1]}"
    local name="${_key_names[$suffix]}"
    if [[ -z "$name" ]]; then
        printf "\033[1;31m_oaki: unrecognized key suffix:\033[0m %s\n" "$suffix" >&2
        return 1
    fi

    printf "%s" "$name"
}

_codex_usage () {
    printf "Usage: %s\n" "$1" >&2
}

_codex_format_key_display () {
    local key="${1:-}"
    local allow_unknown="${2:---allow-unknown}"
    local key_name

    [[ -n "$key" ]] || {
        printf "_codex_format_key_display: key is required\n" >&2
        return 1
    }

    if ! key_name="$(_oaki "$key" 2>/dev/null)"; then
        [[ "$allow_unknown" == "--allow-unknown" ]] || return 1
        key_name="(unknown)"
    fi

    printf "%s (%s…%s)" "$key_name" "${key[1,3]}" "${key[-4,-1]}"
}

_codex_usage_limit_file () {
    printf "%s" "${HOME}/.codex/cul-deadline"
}

_codex_parse_usage_limit_deadline () {
    local message="${1:-}"
    local marker="try again at "
    local deadline_text
    local normalized_deadline
    local deadline_epoch

    [[ -n "$message" ]] || {
        printf "cul: usage-limit message is required\n" >&2
        return 1
    }

    deadline_text="${message#*${marker}}"
    if [[ "$deadline_text" == "$message" ]]; then
        printf "cul: could not find 'try again at ' in the provided message\n" >&2
        return 1
    fi

    deadline_text="${deadline_text%.}"
    normalized_deadline="$(printf "%s\n" "$deadline_text" | sed -E 's/([[:digit:]]+)(st|nd|rd|th)/\1/g')"

    if ! deadline_epoch="$(LC_TIME=C date -j -f "%b %e, %Y %l:%M %p" -v0S "$normalized_deadline" "+%s" 2>/dev/null)"; then
        printf "cul: could not parse usage-limit deadline: %s\n" "$deadline_text" >&2
        return 1
    fi

    printf "%s" "$deadline_epoch"
}

_codex_store_usage_limit_deadline () {
    local deadline_epoch="${1:-}"
    local deadline_file="$(_codex_usage_limit_file)"

    [[ "$deadline_epoch" == <-> ]] || {
        printf "_codex_store_usage_limit_deadline: deadline epoch is required\n" >&2
        return 1
    }

    mkdir -p "${deadline_file:h}" || {
        printf "cul: could not create %s\n" "${deadline_file:h}" >&2
        return 1
    }

    printf "%s\n" "$deadline_epoch" >| "$deadline_file" || {
        printf "cul: could not write %s\n" "$deadline_file" >&2
        return 1
    }
}

_codex_load_usage_limit_deadline () {
    local deadline_file="$(_codex_usage_limit_file)"
    local deadline_epoch

    [[ -r "$deadline_file" ]] || return 1
    deadline_epoch="$(<"$deadline_file")"
    [[ "$deadline_epoch" == <-> ]] || return 1

    printf "%s" "$deadline_epoch"
}

_codex_clear_usage_limit_deadline () {
    local deadline_file="$(_codex_usage_limit_file)"

    [[ -e "$deadline_file" ]] || return 0
    rm -f "$deadline_file" || {
        printf "cul: could not remove %s\n" "$deadline_file" >&2
        return 1
    }
}

_codex_print_usage_limit_guidance () {
    local deadline_epoch="${1:-}"
    local deadline_display
    local now_epoch

    [[ "$deadline_epoch" == <-> ]] || {
        print -P "%F{red}_codex_print_usage_limit_guidance: deadline epoch is required%f" >&2
        return 1
    }

    if ! deadline_display="$(LC_TIME=C date -r "$deadline_epoch" "+%b %e, %Y %l:%M %p %Z" 2>/dev/null)"; then
        print -P "%F{red}_codex_print_usage_limit_guidance: could not format deadline%f" >&2
        return 1
    fi

    now_epoch="$(date "+%s")" || return 1

    if (( now_epoch < deadline_epoch )); then
        print -P "%F{red}Usage-limit deadline: ${deadline_display}%f. Keep using this API key until then, then run '%F{yellow}cul --clear%f'."
        return 0
    fi

    print -P "%F{red}Usage-limit deadline:%f %F{green}${deadline_display} has passed%f. Run '%F{yellow}clo%f' now."
}

_codex_login_with_key () {
    local key="${1:-}"
    local key_display

    [[ -n "$key" ]] || {
        printf "_codex_login_with_key: key is required\n" >&2
        return 1
    }

    setopt localoptions pipefail
    if ! print -rn -- "${key}" | codex login --with-api-key; then
        if key_display="$(_codex_format_key_display "$key")"; then
            printf "Failed to authenticate with codex using %s.\n" "$key_display" >&2
        else
            printf "Failed to authenticate with codex using the provided API key.\n" >&2
        fi
        return 1
    fi

    if ! key_display="$(_codex_format_key_display "$key" --strict)"; then
        printf "\ncli: could not resolve key name; proceeding with masked key only.\n" >&2
        key_display="(unknown) (${key[1,3]}…${key[-4,-1]})"
    fi

    printf "\n\033[2mUsing the following OpenAI API Key:\033[0m \033[1;34m%s\033[0m\n" "$key_display"

    if ! print -rn -- "${key}" | pbcopy; then
        printf "Failed to copy API key to clipboard.\n" >&2
        return 1
    fi

    printf "\n\033[1;31mCAUTION:\033[0m The entire API Key has been copied to the clipboard (so that you can easily login to the Codex Extension).\n"
    printf "\nRemember to 'clo' ([c]odex [l]og[o]ut) at the end of every Visual Studio Code session.\n\n"
}

_codex_read_general_key () {
    local general_key

    printf "Loading %s API Key from 1Password …\n" "${OPENAI_GENERAL_ITEM_NAME}" >&2
    if ! general_key="$(op read "op://Employee/${OPENAI_GENERAL_ITEM_NAME}/password" --no-newline 2>/dev/null)"; then
        printf "Failed to read the %s API Key from 1Password.\n" "${OPENAI_GENERAL_ITEM_NAME}" >&2
        return 1
    fi

    [[ -n "$general_key" ]] || {
        printf "The %s API Key read from 1Password was empty.\n" "${OPENAI_GENERAL_ITEM_NAME}" >&2
        return 1
    }

    printf "%s" "$general_key"
}

_codex_continue_env_file () {
    printf "%s" "${HOME}/.continue/.env"
}

_codex_sync_continue_env_key () {
    local key="${1:-${OPENAI_API_KEY:-}}"
    local env_file="$(_codex_continue_env_file)"
    local env_dir="${env_file:h}"
    local tmp_file="${env_file}.tmp.$$"
    local line
    local found=0

    [[ -n "$key" ]] || {
        printf "_codex_sync_continue_env_key: key is required\n" >&2
        return 1
    }

    mkdir -p "$env_dir" || {
        printf "Failed to create %s.\n" "$env_dir" >&2
        return 1
    }

    : >| "$tmp_file" || {
        printf "Failed to create temporary file for %s.\n" "$env_file" >&2
        return 1
    }

    if [[ -r "$env_file" ]]; then
        while IFS= read -r line || [[ -n "$line" ]]; do
            if [[ "$line" == OPENAI_API_KEY=* ]]; then
                printf "OPENAI_API_KEY=%s\n" "$key" >> "$tmp_file" || return 1
                found=1
            else
                printf "%s\n" "$line" >> "$tmp_file" || return 1
            fi
        done < "$env_file"
    fi

    if (( ! found )); then
        printf "OPENAI_API_KEY=%s\n" "$key" >> "$tmp_file" || return 1
    fi

    mv "$tmp_file" "$env_file" || {
        rm -f "$tmp_file"
        printf "Failed to update %s.\n" "$env_file" >&2
        return 1
    }

    chmod 600 "$env_file" 2>/dev/null || true
}

_codex_clear_continue_env_key () {
    local env_file="$(_codex_continue_env_file)"
    local tmp_file="${env_file}.tmp.$$"
    local line
    local kept=0

    [[ -e "$env_file" ]] || return 0

    : >| "$tmp_file" || {
        printf "Failed to create temporary file for %s.\n" "$env_file" >&2
        return 1
    }

    while IFS= read -r line || [[ -n "$line" ]]; do
        if [[ "$line" == OPENAI_API_KEY=* ]]; then
            continue
        fi
        printf "%s\n" "$line" >> "$tmp_file" || return 1
        kept=1
    done < "$env_file"

    if (( kept )); then
        mv "$tmp_file" "$env_file" || {
            rm -f "$tmp_file"
            printf "Failed to update %s.\n" "$env_file" >&2
            return 1
        }
    else
        rm -f "$tmp_file" "$env_file" || {
            printf "Failed to remove %s.\n" "$env_file" >&2
            return 1
        }
    fi
}

_codex_load_auth_state () {
    local auth_file="${HOME}/.codex/auth.json"
    local auth_json

    typeset -g _CODEX_AUTH_LOADED=0
    typeset -g _CODEX_AUTH_MODE=""
    typeset -g _CODEX_AUTH_KEY=""

    [[ -r "$auth_file" ]] || return 1
    auth_json="$(<"$auth_file")"
    [[ -n "$auth_json" ]] || return 1

    _CODEX_AUTH_LOADED=1
    _CODEX_AUTH_MODE="$(printf "%s\n" "$auth_json" | sed -n 's/^[[:space:]]*"auth_mode"[[:space:]]*:[[:space:]]*"\([^"]*\)".*/\1/p')"
    _CODEX_AUTH_KEY="$(printf "%s\n" "$auth_json" | sed -n 's/^[[:space:]]*"OPENAI_API_KEY"[[:space:]]*:[[:space:]]*"\([^"]*\)".*/\1/p')"
}

_codex_print_shell_env_status () {
    local key="${OPENAI_API_KEY:-}"

    if [[ -z "$key" ]]; then
        printf "\033[2mRepo key loaded in shell:\033[0m none (OPENAI_API_KEY is unset)\n"
        return 0
    fi

    printf "\033[2mRepo key loaded in shell:\033[0m \033[1;34m%s\033[0m\n" "$(_codex_format_key_display "$key")"
}

_codex_print_auth_file_status () {
    if [[ -n "$_CODEX_AUTH_KEY" ]]; then
        if [[ -n "$_CODEX_AUTH_MODE" ]]; then
            printf "\033[2mCodex auth file:\033[0m \033[1;34m%s -> %s\033[0m\n" "$_CODEX_AUTH_MODE" "$(_codex_format_key_display "$_CODEX_AUTH_KEY")"
        else
            printf "\033[2mCodex auth file:\033[0m \033[1;34m%s\033[0m\n" "$(_codex_format_key_display "$_CODEX_AUTH_KEY")"
        fi
        return 0
    fi

    if [[ -n "$_CODEX_AUTH_MODE" ]]; then
        printf "\033[2mCodex auth file:\033[0m \033[1;34m%s (no stored API key)\033[0m\n" "$_CODEX_AUTH_MODE"
        return 0
    fi

    printf "\033[2mCodex auth file:\033[0m \033[1;34msigned out\033[0m\n"
}

_codex_print_state_mismatch_note () {
    local shell_key="${OPENAI_API_KEY:-}"

    [[ "$_CODEX_AUTH_LOADED" -eq 1 ]] || return 0

    if [[ -n "$shell_key" && -z "$_CODEX_AUTH_KEY" ]]; then
        if [[ "$_CODEX_AUTH_MODE" == "chatgpt" ]]; then
            printf "\033[1;31mMismatch:\033[0m a repo key is loaded in the shell, but Codex auth file is chatgpt with no stored API key.\n"
        elif [[ -n "$_CODEX_AUTH_MODE" ]]; then
            printf "\033[1;31mMismatch:\033[0m a repo key is loaded in the shell, but Codex auth file is %s with no stored API key.\n" "$_CODEX_AUTH_MODE"
        else
            printf "\033[1;31mMismatch:\033[0m a repo key is loaded in the shell, but Codex auth file is signed out with no stored API key.\n"
        fi
        return 0
    fi

    if [[ -z "$shell_key" && -n "$_CODEX_AUTH_KEY" ]]; then
        if [[ -n "$_CODEX_AUTH_MODE" ]]; then
            printf "\033[1;31mMismatch:\033[0m no repo key is loaded in the shell, but Codex auth file has a stored API key under %s.\n" "$_CODEX_AUTH_MODE"
        else
            printf "\033[1;31mMismatch:\033[0m no repo key is loaded in the shell, but Codex auth file has a stored API key.\n"
        fi
        return 0
    fi

    if [[ -n "$shell_key" && -n "$_CODEX_AUTH_KEY" && "$shell_key" != "$_CODEX_AUTH_KEY" ]]; then
        printf "\033[1;31mMismatch:\033[0m the repo key loaded in the shell differs from the API key stored in Codex auth.\n"
    fi
}

_codex_print_source_of_truth_note () {
    printf "\033[1;33mNote:\033[0m the shell can load OPENAI_API_KEY independently; Codex persists login state in \033[1;34mjq .OPENAI_API_KEY ~/.codex/auth.json\033[0m.\n"
}

_codex_resolve_login_key () {
    local mode="${1:-}"
    local key

    case "$mode" in
        "")
            if [[ -n "${OPENAI_API_KEY:-}" ]]; then
                printf "%s" "${OPENAI_API_KEY}"
                return 0
            fi

            printf "OPENAI_API_KEY is not set; attempting to load from direnv…\n" >&2
            key="$(direnv exec . sh -c 'printf "%s" "$OPENAI_API_KEY"' 2>/dev/null)"
            if [[ -z "$key" ]]; then
                printf "OPENAI_API_KEY is still not set — is there an .envrc in the current directory?\n" >&2
                return 1
            fi

            printf "%s" "$key"
            ;;
        "--general")
            _codex_read_general_key
            ;;
        *)
            _codex_usage "cli [--general]"
            return 1
            ;;
    esac
}

_codex_logout () {
    if ! codex logout; then
        printf "codex logout failed.\n" >&2
        return 1
    fi
}

cli () {     # [c]odex [l]og[i]n
    local mode="${1:-}"
    local key
    local deadline_epoch

    if ! key="$(_codex_resolve_login_key "$mode")"; then
        return 1
    fi

    if [[ "$mode" != "--general" && -z "${OPENAI_API_KEY:-}" ]]; then
        OPENAI_API_KEY="$key"
        export OPENAI_API_KEY
    fi

    _codex_login_with_key "$key" || return 1
    _codex_sync_continue_env_key "$key" || return 1

    if deadline_epoch="$(_codex_load_usage_limit_deadline)"; then
        _codex_print_usage_limit_guidance "$deadline_epoch"
    fi
}

clo () {     # [c]odex [l]og[o]ut
    if (( $# > 0 )); then
        _codex_usage "clo"
        return 1
    fi

    if [[ -n "${OPENAI_API_KEY:-}" ]]; then
        printf "Clearing OpenAI API Key …\n" >&2
        unset OPENAI_API_KEY || {
            printf "Failed to unset OPENAI_API_KEY.\n" >&2
            return 1
        }
    fi

    _codex_clear_continue_env_key || return 1

    _codex_logout || return 1
    printf "\nRemember to 'Cmd + Shift + P → Developer: Restart Extension Host [ ⌘ ⇧ ⌥ R ]' before logging in again.\n\n"
}

cul () {     # [c]odex [u]sage-[l]imit
    local usage="cul [--clear | '<usage-limit-message>']"
    local deadline_epoch

    (( $# == 1 )) || {
        _codex_usage "$usage"
        return 1
    }

    if [[ "$1" == "--clear" ]]; then
        _codex_clear_usage_limit_deadline || return 1
        printf "Cleared stored usage-limit reminder.\n"
        return 0
    fi

    if ! deadline_epoch="$(_codex_parse_usage_limit_deadline "$1")"; then
        return 1
    fi

    _codex_store_usage_limit_deadline "$deadline_epoch" || return 1
    _codex_print_usage_limit_guidance "$deadline_epoch"
}

cra () {     # [c]odex [r]e-[a]uthenticate
    _codex_logout || return 1
    cli
}

cai () {     # [c]odex [a]pi [i]nfo
    local usage="cai [--shell | --codex | <api-key>]"
    local key

    case "${1:-}" in
        "")
            (( $# == 0 )) || {
                _codex_usage "$usage"
                return 1
            }

            _codex_load_auth_state 2>/dev/null || true
            _codex_print_shell_env_status
            _codex_print_auth_file_status
            _codex_print_state_mismatch_note
            _codex_print_source_of_truth_note
            return 0
            ;;
        "--shell")
            (( $# == 1 )) || {
                _codex_usage "$usage"
                return 1
            }

            _codex_print_shell_env_status
            return 0
            ;;
        "--codex")
            (( $# == 1 )) || {
                _codex_usage "$usage"
                return 1
            }

            _codex_load_auth_state 2>/dev/null || true
            _codex_print_auth_file_status
            return 0
            ;;
        --*)
            _codex_usage "$usage"
            return 1
            ;;
    esac

    (( $# == 1 )) || {
        _codex_usage "$usage"
        return 1
    }

    key="$1"
    [[ -n "$key" ]] || {
        printf "cai: key is required\n" >&2
        return 1
    }

    if ! key="$(_codex_format_key_display "$key")"; then
        printf "cai: could not format the provided OpenAI API key\n" >&2
        return 1
    fi

    printf "OpenAI API Key: %s\n" "$key"
}

# --- OpenAI Codex: End -------------------------------------------------------