From 835eb77a7dc1cbb5595cbde793b56eee58382f3e Mon Sep 17 00:00:00 2001
From: Fouad Matin <169186268+fouad-openai@users.noreply.github.com>
Date: Sat, 17 May 2025 21:27:02 -0700
Subject: [PATCH] fix: persist token after refresh (#1006)

After a token refresh/exchange, persist the new refresh and id token
---
 codex-cli/src/cli.tsx               | 19 +++++++++++-------
 codex-cli/src/utils/get-api-key.tsx | 31 ++++++++++++++++++++++++++---
 2 files changed, 40 insertions(+), 10 deletions(-)

diff --git a/codex-cli/src/cli.tsx b/codex-cli/src/cli.tsx
index 7289057f..dfa2b3bb 100644
--- a/codex-cli/src/cli.tsx
+++ b/codex-cli/src/cli.tsx
@@ -333,15 +333,20 @@ if (cli.flags.login) {
 // Ensure the API key is available as an environment variable for legacy code
 process.env["OPENAI_API_KEY"] = apiKey;
 
-if (cli.flags.free && savedTokens?.refresh_token) {
+if (cli.flags.free) {
   // eslint-disable-next-line no-console
   console.log(`${chalk.bold("codex --free")} attempting to redeem credits...`);
-  await maybeRedeemCredits(
-    client.issuer,
-    client.client_id,
-    savedTokens.refresh_token,
-    savedTokens.id_token,
-  );
+  if (!savedTokens?.refresh_token) {
+    apiKey = await fetchApiKey(client.issuer, client.client_id, true);
+    // fetchApiKey includes credit redemption as the end of the flow
+  } else {
+    await maybeRedeemCredits(
+      client.issuer,
+      client.client_id,
+      savedTokens.refresh_token,
+      savedTokens.id_token,
+    );
+  }
 }
 
 // Set of providers that don't require API keys
diff --git a/codex-cli/src/utils/get-api-key.tsx b/codex-cli/src/utils/get-api-key.tsx
index c35dd2a5..4817e396 100644
--- a/codex-cli/src/utils/get-api-key.tsx
+++ b/codex-cli/src/utils/get-api-key.tsx
@@ -132,13 +132,37 @@ async function maybeRedeemCredits(
           );
           return;
         }
-        const refreshData = (await refreshRes.json()) as { id_token: string };
+        const refreshData = (await refreshRes.json()) as {
+          id_token: string;
+          refresh_token?: string;
+        };
         currentIdToken = refreshData.id_token;
         idClaims = JSON.parse(
           Buffer.from(currentIdToken.split(".")[1]!, "base64url").toString(
             "utf8",
           ),
         ) as IDTokenClaims;
+        if (refreshData.refresh_token) {
+          try {
+            const home = os.homedir();
+            const authDir = path.join(home, ".codex");
+            const authFile = path.join(authDir, "auth.json");
+            const existingJson = JSON.parse(
+              await fs.readFile(authFile, "utf-8"),
+            );
+            existingJson.tokens.id_token = currentIdToken;
+            existingJson.tokens.refresh_token = refreshData.refresh_token;
+            existingJson.last_refresh = new Date().toISOString();
+            await fs.writeFile(
+              authFile,
+              JSON.stringify(existingJson, null, 2),
+              { mode: 0o600 },
+            );
+          } catch (err) {
+            // eslint-disable-next-line no-console
+            console.warn("Unable to update refresh token in auth file:", err);
+          }
+        }
       } catch (err) {
         // eslint-disable-next-line no-console
         console.warn("Unable to refresh ID token via token-exchange:", err);
@@ -200,7 +224,7 @@ async function maybeRedeemCredits(
     const redeemRes = await fetch(`${apiHost}/v1/billing/redeem_credits`, {
       method: "POST",
       headers: { "Content-Type": "application/json" },
-      body: JSON.stringify({ id_token: idToken }),
+      body: JSON.stringify({ id_token: currentIdToken }),
     });
 
     if (!redeemRes.ok) {
@@ -713,8 +737,9 @@ async function signInFlow(issuer: string, clientId: string): Promise<string> {
 export async function getApiKey(
   issuer: string,
   clientId: string,
+  forceLogin: boolean = false,
 ): Promise<string> {
-  if (process.env["OPENAI_API_KEY"]) {
+  if (!forceLogin && process.env["OPENAI_API_KEY"]) {
     return process.env["OPENAI_API_KEY"]!;
   }
   const choice = await promptUserForChoice();