From e59083f73bbaa9fc3331dd3fce2f5edd703034fe Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Wed, 4 Jun 2025 21:23:26 -0700 Subject: [PATCH] feat: port maybeRedeemCredits() from get-api-key.tsx to login_with_chatgpt.py --- codex-rs/login/src/login_with_chatgpt.py | 202 ++++++++++++++++++++++- 1 file changed, 201 insertions(+), 1 deletion(-) diff --git a/codex-rs/login/src/login_with_chatgpt.py b/codex-rs/login/src/login_with_chatgpt.py index c1d478644b..31c3448e88 100644 --- a/codex-rs/login/src/login_with_chatgpt.py +++ b/codex-rs/login/src/login_with_chatgpt.py @@ -27,6 +27,7 @@ import urllib.parse import urllib.request import webbrowser from dataclasses import dataclass +from typing import Any, Dict # for type hints # Required port for OAuth client. REQUIRED_PORT = 1455 @@ -313,7 +314,20 @@ class _ApiKeyHTTPHandler(http.server.BaseHTTPRequestHandler): } success_url = f"{URL_BASE}/success?{urllib.parse.urlencode(success_url_query)}" - # TODO(mbolin): Port maybeRedeemCredits() to Python and call it here. + # Attempt to redeem complimentary API credits for eligible ChatGPT + # Plus / Pro subscribers. Any errors are logged but do not interrupt + # the login flow. + + try: + maybe_redeem_credits( + issuer=self.server.issuer, + client_id=self.server.client_id, + refresh_token=token_data.refresh_token, + id_token=token_data.id_token, + codex_home=self.server.codex_home, + ) + except Exception as exc: # pragma: no cover – best-effort only + eprint(f"Unable to redeem ChatGPT subscriber API credits: {exc}") # Persist refresh_token/id_token for future use (redeem credits etc.) last_refresh_str = ( @@ -429,6 +443,192 @@ def eprint(*args, **kwargs) -> None: print(*args, file=sys.stderr, **kwargs) +# --------------------------------------------------------------------------- +# Credit redemption helper – Python port of codex-cli/src/utils/get-api-key.tsx +# --------------------------------------------------------------------------- + + +def _decode_jwt_segment(segment: str) -> Dict[str, Any]: + """Return the decoded JSON payload from a JWT segment. + + Adds required padding for urlsafe_b64decode. + """ + + padded = segment + "=" * (-len(segment) % 4) + try: + data = base64.urlsafe_b64decode(padded.encode()) + return json.loads(data.decode()) + except Exception: + return {} + + +def _current_timestamp_ms() -> int: + return int(datetime.datetime.now(datetime.timezone.utc).timestamp() * 1000) + + +def maybe_redeem_credits( + *, + issuer: str, + client_id: str, + refresh_token: str, + id_token: str | None, + codex_home: str, +) -> None: + """Attempt to redeem complimentary API credits for ChatGPT subscribers. + + The operation is best-effort: any error results in a warning being printed + and the function returning early without raising. + """ + + try: + current_id_token = id_token or "" + + # ----------------------------------------------------------------- + # Parse ID-token claims (if provided) + # ----------------------------------------------------------------- + id_claims: Dict[str, Any] | None = None + if current_id_token and "." in current_id_token: + parts = current_id_token.split(".") + if len(parts) >= 2: + id_claims = _decode_jwt_segment(parts[1]) # type: ignore[arg-type] + + # ----------------------------------------------------------------- + # Refresh expired ID token, if possible + # ----------------------------------------------------------------- + token_expired = True + if id_claims and isinstance(id_claims.get("exp"), (int, float)): + token_expired = _current_timestamp_ms() >= int(id_claims["exp"]) * 1000 + + if token_expired: + eprint("Refreshing credentials...") + try: + payload = json.dumps( + { + "client_id": client_id, + "grant_type": "refresh_token", + "refresh_token": refresh_token, + "scope": "openid profile email", + } + ).encode() + + req = urllib.request.Request( + url="https://auth.openai.com/oauth/token", + data=payload, + method="POST", + headers={"Content-Type": "application/json"}, + ) + + with urllib.request.urlopen(req) as resp: + refresh_data = json.loads(resp.read().decode()) + + current_id_token = refresh_data.get("id_token", current_id_token) + id_claims = None + if current_id_token and "." in current_id_token: + parts = current_id_token.split(".") + if len(parts) >= 2: + id_claims = _decode_jwt_segment(parts[1]) + + new_refresh_token = refresh_data.get("refresh_token") + + # Update auth.json with new tokens, if we have one + if new_refresh_token or current_id_token: + try: + auth_dir = codex_home + auth_path = os.path.join(auth_dir, "auth.json") + with open(auth_path, "r", encoding="utf-8") as fp: + existing = json.load(fp) + + if current_id_token: + existing.setdefault("tokens", {})["id_token"] = current_id_token + if new_refresh_token: + existing.setdefault("tokens", {})["refresh_token"] = new_refresh_token + existing["last_refresh"] = ( + datetime.datetime.now(datetime.timezone.utc) + .isoformat() + .replace("+00:00", "Z") + ) + + with open(auth_path, "w", encoding="utf-8") as fp: + if hasattr(os, "fchmod"): + os.fchmod(fp.fileno(), 0o600) + json.dump(existing, fp, indent=2) + except Exception as err: + eprint("Unable to update refresh token in auth file:", err) + except Exception as err: + eprint("Unable to refresh ID token via token-exchange:", err) + return + + if not id_claims: + # Still couldn't parse claims + return + + auth_claims = id_claims.get("https://api.openai.com/auth", {}) + + # ----------------------------------------------------------------- + # Subscription eligibility check (Plus or Pro, >7 days active) + # ----------------------------------------------------------------- + sub_start_str = auth_claims.get("chatgpt_subscription_active_start") + if isinstance(sub_start_str, str): + try: + sub_start_ts = datetime.datetime.fromisoformat(sub_start_str.rstrip("Z")) + if ( + datetime.datetime.now(datetime.timezone.utc) - sub_start_ts + < datetime.timedelta(days=7) + ): + eprint( + "Sorry, your subscription must be active for more than 7 days to redeem credits." + ) + return + except ValueError: + # Malformed; ignore + pass + + completed_onboarding = bool(auth_claims.get("completed_platform_onboarding")) + is_org_owner = bool(auth_claims.get("is_org_owner")) + needs_setup = not completed_onboarding and is_org_owner + + plan_type = auth_claims.get("chatgpt_plan_type") + + if needs_setup or plan_type not in {"plus", "pro"}: + eprint( + "Users with Plus or Pro subscriptions can redeem free API credits." + ) + return + + api_host = ( + "https://api.openai.com" if issuer == "https://auth.openai.com" else "https://api.openai.org" + ) + + try: + redeem_payload = json.dumps({"id_token": current_id_token}).encode() + req = urllib.request.Request( + url=f"{api_host}/v1/billing/redeem_credits", + data=redeem_payload, + method="POST", + headers={"Content-Type": "application/json"}, + ) + + with urllib.request.urlopen(req) as resp: + redeem_data = json.loads(resp.read().decode()) + + granted = redeem_data.get("granted_chatgpt_subscriber_api_credits", 0) + if granted and granted > 0: + print( + f"Thanks for being a ChatGPT {'Plus' if plan_type=='plus' else 'Pro'} subscriber! " + f"If you haven't already redeemed, you should receive {'$5' if plan_type=='plus' else '$50'} in API credits.", + file=sys.stderr, + ) + else: + eprint("It looks like no credits were granted:") + eprint(json.dumps(redeem_data, indent=2)) + except Exception as err: + eprint("Credit redemption request failed:", err) + + except Exception as exc: + eprint("Unable to redeem ChatGPT subscriber API credits:", exc) + + + LOGIN_SUCCESS_HTML = """