248 lines
8.2 KiB
TypeScript
248 lines
8.2 KiB
TypeScript
import { describe, expect, test } from "bun:test"
|
|
import {
|
|
CodexAuthPlugin,
|
|
parseJwtClaims,
|
|
extractAccountIdFromClaims,
|
|
extractAccountId,
|
|
type IdTokenClaims,
|
|
} from "../../src/plugin/openai/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")
|
|
})
|
|
})
|
|
|
|
test("installs websocket transport only when experimental websockets are enabled", async () => {
|
|
const disabled = await CodexAuthPlugin({} as never)
|
|
const enabled = await CodexAuthPlugin({} as never, { experimentalWebSockets: true })
|
|
|
|
const disabledOptions = await disabled.auth!.loader!(
|
|
async () => ({ type: "api", key: "sk-test" }) as never,
|
|
{} as never,
|
|
)
|
|
const enabledOptions = await enabled.auth!.loader!(
|
|
async () => ({ type: "api", key: "sk-test" }) as never,
|
|
{} as never,
|
|
)
|
|
|
|
expect(disabledOptions.fetch).toBeUndefined()
|
|
expect(enabledOptions.fetch).toBeFunction()
|
|
await enabled.dispose?.()
|
|
})
|
|
|
|
test("deduplicates concurrent Codex token refreshes", async () => {
|
|
let auth = {
|
|
type: "oauth" as const,
|
|
refresh: "refresh-old",
|
|
access: "",
|
|
expires: 0,
|
|
}
|
|
const authUpdates: Array<{
|
|
body: { refresh: string; access: string; expires: number; accountId?: string }
|
|
}> = []
|
|
let resolveRefresh: (() => void) | undefined
|
|
const refreshReady = new Promise<void>((resolve) => {
|
|
resolveRefresh = resolve
|
|
})
|
|
let refreshRequests = 0
|
|
const apiRequests: { authorization: string | null; accountId: string | null }[] = []
|
|
|
|
using server = Bun.serve({
|
|
port: 0,
|
|
async fetch(request) {
|
|
const url = new URL(request.url)
|
|
if (url.pathname === "/oauth/token") {
|
|
expect(await request.text()).toContain("refresh_token=refresh-old")
|
|
refreshRequests += 1
|
|
await refreshReady
|
|
return Response.json({
|
|
id_token: createTestJwt({ chatgpt_account_id: "acc-123" }),
|
|
access_token: "access-new",
|
|
refresh_token: "refresh-new",
|
|
expires_in: 3600,
|
|
})
|
|
}
|
|
|
|
if (url.pathname === "/backend-api/codex/responses") {
|
|
apiRequests.push({
|
|
authorization: request.headers.get("authorization"),
|
|
accountId: request.headers.get("ChatGPT-Account-Id"),
|
|
})
|
|
return new Response("{}", { status: 200 })
|
|
}
|
|
|
|
return new Response("unexpected request", { status: 500 })
|
|
},
|
|
})
|
|
|
|
const hooks = await CodexAuthPlugin(
|
|
{
|
|
client: {
|
|
auth: {
|
|
async set(input: { body: { refresh: string; access: string; expires: number; accountId?: string } }) {
|
|
authUpdates.push(input)
|
|
auth = {
|
|
type: "oauth",
|
|
refresh: input.body.refresh,
|
|
access: input.body.access,
|
|
expires: input.body.expires,
|
|
...(input.body.accountId && { accountId: input.body.accountId }),
|
|
}
|
|
},
|
|
},
|
|
} as never,
|
|
project: {} as never,
|
|
directory: "",
|
|
worktree: "",
|
|
experimental_workspace: {
|
|
register() {},
|
|
},
|
|
serverUrl: new URL("https://example.com"),
|
|
$: {} as never,
|
|
},
|
|
{
|
|
issuer: server.url.origin,
|
|
codexApiEndpoint: new URL("/backend-api/codex/responses", server.url).toString(),
|
|
},
|
|
)
|
|
const loaded = await hooks.auth!.loader!(async () => auth as never, {} as never)
|
|
|
|
const first = loaded.fetch!("https://api.openai.com/v1/responses")
|
|
const second = loaded.fetch!("https://api.openai.com/v1/responses")
|
|
|
|
await waitFor(() => refreshRequests === 1)
|
|
expect(apiRequests).toHaveLength(0)
|
|
|
|
resolveRefresh!()
|
|
await Promise.all([first, second])
|
|
|
|
expect(refreshRequests).toBe(1)
|
|
expect(authUpdates).toHaveLength(1)
|
|
expect(authUpdates[0]?.body.refresh).toBe("refresh-new")
|
|
expect(authUpdates[0]?.body.access).toBe("access-new")
|
|
expect(authUpdates[0]?.body.accountId).toBe("acc-123")
|
|
expect(apiRequests).toEqual([
|
|
{ authorization: "Bearer access-new", accountId: "acc-123" },
|
|
{ authorization: "Bearer access-new", accountId: "acc-123" },
|
|
])
|
|
})
|
|
})
|
|
|
|
async function waitFor(predicate: () => boolean) {
|
|
const started = Date.now()
|
|
while (!predicate()) {
|
|
if (Date.now() - started > 1_000) throw new Error("timed out waiting for condition")
|
|
await new Promise((resolve) => setTimeout(resolve, 1))
|
|
}
|
|
}
|