fix: add ChatGPT-Account-Id header for organization subscriptions (#7603)

4 files changed, 186 insertions, 12 deletions

diff --git a/packages/opencode/src/auth/index.ts b/packages/opencode/src/auth/index.ts
index 6642a0742..3fd283053 100644
--- a/packages/opencode/src/auth/index.ts
+++ b/packages/opencode/src/auth/index.ts
@@ -12,6 +12,7 @@ export namespace Auth {
       refresh: z.string(),
       access: z.string(),
       expires: z.number(),
+      accountId: z.string().optional(),
       enterpriseUrl: z.string().optional(),
     })
     .meta({ ref: "OAuth" })
diff --git a/packages/opencode/src/plugin/codex.ts b/packages/opencode/src/plugin/codex.ts
index d0f025b06..4e2b28379 100644
--- a/packages/opencode/src/plugin/codex.ts
+++ b/packages/opencode/src/plugin/codex.ts
@@ -42,6 +42,46 @@ function generateState(): string {
   return base64UrlEncode(crypto.getRandomValues(new Uint8Array(32)).buffer)
 }
 
+export interface IdTokenClaims {
+  chatgpt_account_id?: string
+  organizations?: Array<{ id: string }>
+  email?: string
+  "https://api.openai.com/auth"?: {
+    chatgpt_account_id?: string
+  }
+}
+
+export function parseJwtClaims(token: string): IdTokenClaims | undefined {
+  const parts = token.split(".")
+  if (parts.length !== 3) return undefined
+  try {
+    return JSON.parse(Buffer.from(parts[1], "base64url").toString())
+  } catch {
+    return undefined
+  }
+}
+
+export function extractAccountIdFromClaims(claims: IdTokenClaims): string | undefined {
+  return (
+    claims.chatgpt_account_id ||
+    claims["https://api.openai.com/auth"]?.chatgpt_account_id ||
+    claims.organizations?.[0]?.id
+  )
+}
+
+export function extractAccountId(tokens: TokenResponse): string | undefined {
+  if (tokens.id_token) {
+    const claims = parseJwtClaims(tokens.id_token)
+    const accountId = claims && extractAccountIdFromClaims(claims)
+    if (accountId) return accountId
+  }
+  if (tokens.access_token) {
+    const claims = parseJwtClaims(tokens.access_token)
+    return claims ? extractAccountIdFromClaims(claims) : undefined
+  }
+  return undefined
+}
+
 function buildAuthorizeUrl(redirectUri: string, pkce: PkceCodes, state: string): string {
   const params = new URLSearchParams({
     response_type: "code",
@@ -380,10 +420,14 @@ export async function CodexAuthPlugin(input: PluginInput): Promise<Hooks> {
             const currentAuth = await getAuth()
             if (currentAuth.type !== "oauth") return fetch(requestInput, init)
 
+            // Cast to include accountId field
+            const authWithAccount = currentAuth as typeof currentAuth & { accountId?: string }
+
             // Check if token needs refresh
             if (!currentAuth.access || currentAuth.expires < Date.now()) {
               log.info("refreshing codex access token")
               const tokens = await refreshAccessToken(currentAuth.refresh)
+              const newAccountId = extractAccountId(tokens) || authWithAccount.accountId
               await input.client.auth.set({
                 path: { id: "codex" },
                 body: {
@@ -391,9 +435,11 @@ export async function CodexAuthPlugin(input: PluginInput): Promise<Hooks> {
                   refresh: tokens.refresh_token,
                   access: tokens.access_token,
                   expires: Date.now() + (tokens.expires_in ?? 3600) * 1000,
+                  ...(newAccountId && { accountId: newAccountId }),
                 },
               })
               currentAuth.access = tokens.access_token
+              authWithAccount.accountId = newAccountId
             }
 
             // Build headers
@@ -415,20 +461,20 @@ export async function CodexAuthPlugin(input: PluginInput): Promise<Hooks> {
             // Set authorization header with access token
             headers.set("authorization", `Bearer ${currentAuth.access}`)
 
-            // Rewrite URL to Codex endpoint
-            let url: URL
-            if (typeof requestInput === "string") {
-              url = new URL(requestInput)
-            } else if (requestInput instanceof URL) {
-              url = requestInput
-            } else {
-              url = new URL(requestInput.url)
+            // Set ChatGPT-Account-Id header for organization subscriptions
+            if (authWithAccount.accountId) {
+              headers.set("ChatGPT-Account-Id", authWithAccount.accountId)
             }
 
-            // If this is a messages/responses request, redirect to Codex endpoint
-            if (url.pathname.includes("/v1/responses") || url.pathname.includes("/chat/completions")) {
-              url = new URL(CODEX_API_ENDPOINT)
-            }
+            // Rewrite URL to Codex endpoint
+            const parsed =
+              requestInput instanceof URL
+                ? requestInput
+                : new URL(typeof requestInput === "string" ? requestInput : requestInput.url)
+            const url =
+              parsed.pathname.includes("/v1/responses") || parsed.pathname.includes("/chat/completions")
+                ? new URL(CODEX_API_ENDPOINT)
+                : parsed
 
             return fetch(url, {
               ...init,
@@ -456,11 +502,13 @@ export async function CodexAuthPlugin(input: PluginInput): Promise<Hooks> {
               callback: async () => {
                 const tokens = await callbackPromise
                 stopOAuthServer()
+                const accountId = extractAccountId(tokens)
                 return {
                   type: "success" as const,
                   refresh: tokens.refresh_token,
                   access: tokens.access_token,
                   expires: Date.now() + (tokens.expires_in ?? 3600) * 1000,
+                  accountId,
                 }
               },
             }
diff --git a/packages/opencode/test/plugin/codex.test.ts b/packages/opencode/test/plugin/codex.test.ts
new file mode 100644
index 000000000..74d28ac9d
--- /dev/null
+++ b/packages/opencode/test/plugin/codex.test.ts
@@ -0,0 +1,123 @@
+import { describe, expect, test } from "bun:test"
+import {
+  parseJwtClaims,
+  extractAccountIdFromClaims,
+  extractAccountId,
+  type IdTokenClaims,
+} from "../../src/plugin/codex"
+
+function createTestJwt(payload: object): string {
+  const header = Buffer.from(JSON.stringify({ alg: "none" })).toString("base64url")
+  const body = Buffer.from(JSON.stringify(payload)).toString("base64url")
+  return `${header}.${body}.sig`
+}
+
+describe("plugin.codex", () => {
+  describe("parseJwtClaims", () => {
+    test("parses valid JWT with claims", () => {
+      const payload = { email: "[email protected]", chatgpt_account_id: "acc-123" }
+      const jwt = createTestJwt(payload)
+      const claims = parseJwtClaims(jwt)
+      expect(claims).toEqual(payload)
+    })
+
+    test("returns undefined for JWT with less than 3 parts", () => {
+      expect(parseJwtClaims("invalid")).toBeUndefined()
+      expect(parseJwtClaims("only.two")).toBeUndefined()
+    })
+
+    test("returns undefined for invalid base64", () => {
+      expect(parseJwtClaims("a.!!!invalid!!!.b")).toBeUndefined()
+    })
+
+    test("returns undefined for invalid JSON payload", () => {
+      const header = Buffer.from("{}").toString("base64url")
+      const invalidJson = Buffer.from("not json").toString("base64url")
+      expect(parseJwtClaims(`${header}.${invalidJson}.sig`)).toBeUndefined()
+    })
+  })
+
+  describe("extractAccountIdFromClaims", () => {
+    test("extracts chatgpt_account_id from root", () => {
+      const claims: IdTokenClaims = { chatgpt_account_id: "acc-root" }
+      expect(extractAccountIdFromClaims(claims)).toBe("acc-root")
+    })
+
+    test("extracts chatgpt_account_id from nested https://api.openai.com/auth", () => {
+      const claims: IdTokenClaims = {
+        "https://api.openai.com/auth": { chatgpt_account_id: "acc-nested" },
+      }
+      expect(extractAccountIdFromClaims(claims)).toBe("acc-nested")
+    })
+
+    test("prefers root over nested", () => {
+      const claims: IdTokenClaims = {
+        chatgpt_account_id: "acc-root",
+        "https://api.openai.com/auth": { chatgpt_account_id: "acc-nested" },
+      }
+      expect(extractAccountIdFromClaims(claims)).toBe("acc-root")
+    })
+
+    test("extracts from organizations array as fallback", () => {
+      const claims: IdTokenClaims = {
+        organizations: [{ id: "org-123" }, { id: "org-456" }],
+      }
+      expect(extractAccountIdFromClaims(claims)).toBe("org-123")
+    })
+
+    test("returns undefined when no accountId found", () => {
+      const claims: IdTokenClaims = { email: "[email protected]" }
+      expect(extractAccountIdFromClaims(claims)).toBeUndefined()
+    })
+  })
+
+  describe("extractAccountId", () => {
+    test("extracts from id_token first", () => {
+      const idToken = createTestJwt({ chatgpt_account_id: "from-id-token" })
+      const accessToken = createTestJwt({ chatgpt_account_id: "from-access-token" })
+      expect(
+        extractAccountId({
+          id_token: idToken,
+          access_token: accessToken,
+          refresh_token: "rt",
+        }),
+      ).toBe("from-id-token")
+    })
+
+    test("falls back to access_token when id_token has no accountId", () => {
+      const idToken = createTestJwt({ email: "[email protected]" })
+      const accessToken = createTestJwt({
+        "https://api.openai.com/auth": { chatgpt_account_id: "from-access" },
+      })
+      expect(
+        extractAccountId({
+          id_token: idToken,
+          access_token: accessToken,
+          refresh_token: "rt",
+        }),
+      ).toBe("from-access")
+    })
+
+    test("returns undefined when no tokens have accountId", () => {
+      const token = createTestJwt({ email: "[email protected]" })
+      expect(
+        extractAccountId({
+          id_token: token,
+          access_token: token,
+          refresh_token: "rt",
+        }),
+      ).toBeUndefined()
+    })
+
+    test("handles missing id_token", () => {
+      const accessToken = createTestJwt({ chatgpt_account_id: "acc-123" })
+      expect(
+        extractAccountId({
+          id_token: "",
+          access_token: accessToken,
+          refresh_token: "rt",
+        }),
+      ).toBe("acc-123")
+    })
+  })
+})
diff --git a/packages/plugin/src/index.ts b/packages/plugin/src/index.ts
index bf9b6e8c2..46ad8512c 100644
--- a/packages/plugin/src/index.ts
+++ b/packages/plugin/src/index.ts
@@ -114,6 +114,7 @@ export type AuthOuathResult = { url: string; instructions: string } & (
                 refresh: string
                 access: string
                 expires: number
+                accountId?: string
               }
             | { key: string }
           ))
@@ -133,6 +134,7 @@ export type AuthOuathResult = { url: string; instructions: string } & (
                 refresh: string
                 access: string
                 expires: number
+                accountId?: string
               }
             | { key: string }
           ))