diff --git a/packages/opencode/src/auth/index.ts b/packages/opencode/src/auth/index.ts index 6642a07429d..3fd28305368 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 d0f025b0614..4e2b283795d 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 { 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 { 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 { // 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 { 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 00000000000..74d28ac9dcc --- /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: "test@example.com", 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: "test@example.com" } + 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: "test@example.com" }) + 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: "test@example.com" }) + 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 bf9b6e8c2d8..46ad8512c68 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 } ))