openAgent

1import { describe, expect, test } from "bun:test"
2import {
3  CodexAuthPlugin,
4  parseJwtClaims,
5  extractAccountIdFromClaims,
6  extractAccountId,
7  renderOAuthError,
8  type IdTokenClaims,
9} from "../../src/plugin/openai/codex"
10 
11function createTestJwt(payload: object): string {
12  const header = Buffer.from(JSON.stringify({ alg: "none" })).toString("base64url")
13  const body = Buffer.from(JSON.stringify(payload)).toString("base64url")
14  return `${header}.${body}.sig`
15}
16 
17describe("plugin.codex", () => {
18  test("escapes provider errors in callback HTML", () => {
19    const error = `</div><script>alert("xss" & 'more')</script>`
20    const html = renderOAuthError(error)
21 
22    expect(html).toContain("&lt;/div&gt;&lt;script&gt;alert(&quot;xss&quot; &amp; &#39;more&#39;)&lt;/script&gt;")
23    expect(html).not.toContain(error)
24  })
25 
26  describe("parseJwtClaims", () => {
27    test("parses valid JWT with claims", () => {
28      const payload = { email: "test@example.com", chatgpt_account_id: "acc-123" }
29      const jwt = createTestJwt(payload)
30      const claims = parseJwtClaims(jwt)
31      expect(claims).toEqual(payload)
32    })
33 
34    test("returns undefined for JWT with less than 3 parts", () => {
35      expect(parseJwtClaims("invalid")).toBeUndefined()
36      expect(parseJwtClaims("only.two")).toBeUndefined()
37    })
38 
39    test("returns undefined for invalid base64", () => {
40      expect(parseJwtClaims("a.!!!invalid!!!.b")).toBeUndefined()
41    })
42 
43    test("returns undefined for invalid JSON payload", () => {
44      const header = Buffer.from("{}").toString("base64url")
45      const invalidJson = Buffer.from("not json").toString("base64url")
46      expect(parseJwtClaims(`${header}.${invalidJson}.sig`)).toBeUndefined()
47    })
48  })
49 
50  describe("extractAccountIdFromClaims", () => {
51    test("extracts chatgpt_account_id from root", () => {
52      const claims: IdTokenClaims = { chatgpt_account_id: "acc-root" }
53      expect(extractAccountIdFromClaims(claims)).toBe("acc-root")
54    })
55 
56    test("extracts chatgpt_account_id from nested https://api.openai.com/auth", () => {
57      const claims: IdTokenClaims = {
58        "https://api.openai.com/auth": { chatgpt_account_id: "acc-nested" },
59      }
60      expect(extractAccountIdFromClaims(claims)).toBe("acc-nested")
61    })
62 
63    test("prefers root over nested", () => {
64      const claims: IdTokenClaims = {
65        chatgpt_account_id: "acc-root",
66        "https://api.openai.com/auth": { chatgpt_account_id: "acc-nested" },
67      }
68      expect(extractAccountIdFromClaims(claims)).toBe("acc-root")
69    })
70 
71    test("extracts from organizations array as fallback", () => {
72      const claims: IdTokenClaims = {
73        organizations: [{ id: "org-123" }, { id: "org-456" }],
74      }
75      expect(extractAccountIdFromClaims(claims)).toBe("org-123")
76    })
77 
78    test("returns undefined when no accountId found", () => {
79      const claims: IdTokenClaims = { email: "test@example.com" }
80      expect(extractAccountIdFromClaims(claims)).toBeUndefined()
81    })
82  })
83 
84  describe("extractAccountId", () => {
85    test("extracts from id_token first", () => {
86      const idToken = createTestJwt({ chatgpt_account_id: "from-id-token" })
87      const accessToken = createTestJwt({ chatgpt_account_id: "from-access-token" })
88      expect(
89        extractAccountId({
90          id_token: idToken,
91          access_token: accessToken,
92          refresh_token: "rt",
93        }),
94      ).toBe("from-id-token")
95    })
96 
97    test("falls back to access_token when id_token has no accountId", () => {
98      const idToken = createTestJwt({ email: "test@example.com" })
99      const accessToken = createTestJwt({
100        "https://api.openai.com/auth": { chatgpt_account_id: "from-access" },
101      })
102      expect(
103        extractAccountId({
104          id_token: idToken,
105          access_token: accessToken,
106          refresh_token: "rt",
107        }),
108      ).toBe("from-access")
109    })
110 
111    test("returns undefined when no tokens have accountId", () => {
112      const token = createTestJwt({ email: "test@example.com" })
113      expect(
114        extractAccountId({
115          id_token: token,
116          access_token: token,
117          refresh_token: "rt",
118        }),
119      ).toBeUndefined()
120    })
121 
122    test("handles missing id_token", () => {
123      const accessToken = createTestJwt({ chatgpt_account_id: "acc-123" })
124      expect(
125        extractAccountId({
126          id_token: "",
127          access_token: accessToken,
128          refresh_token: "rt",
129        }),
130      ).toBe("acc-123")
131    })
132  })
133 
134  test("installs websocket transport only when experimental websockets are enabled", async () => {
135    const disabled = await CodexAuthPlugin({} as never)
136    const enabled = await CodexAuthPlugin({} as never, { experimentalWebSockets: true })
137 
138    const disabledOptions = await disabled.auth!.loader!(
139      async () => ({ type: "api", key: "sk-test" }) as never,
140      {} as never,
141    )
142    const enabledOptions = await enabled.auth!.loader!(
143      async () => ({ type: "api", key: "sk-test" }) as never,
144      {} as never,
145    )
146 
147    expect(disabledOptions.fetch).toBeUndefined()
148    expect(enabledOptions.fetch).toBeFunction()
149    await enabled.dispose?.()
150  })
151 
152  test("filters unsupported modes and uses Codex context limits for OAuth GPT models", async () => {
153    const hooks = await CodexAuthPlugin({} as never)
154    const limit = { context: 1_050_000, input: 922_000, output: 128_000 }
155    const provider = {
156      models: {
157        ...Object.fromEntries(
158          ["gpt-5.4", "gpt-5.5", "gpt-5.6-sol", "gpt-5.6-terra", "gpt-5.6-luna", "gpt-5.7-pro"].map((id) => [
159            id,
160            { id, api: { id }, limit, cost: {}, options: {} },
161          ]),
162        ),
163        "gpt-5.4-pro": {
164          id: "gpt-5.4-pro",
165          api: { id: "gpt-5.4" },
166          limit,
167          cost: {},
168          options: { reasoningMode: "pro" },
169        },
170        "gpt-5.6-sol-high": {
171          id: "gpt-5.6-sol-high",
172          api: { id: "gpt-5.6-sol" },
173          limit,
174          cost: {},
175          options: { reasoningEffort: "high" },
176        },
177      },
178    }
179 
180    const models = await hooks.provider!.models!(provider as never, { auth: { type: "oauth" } } as never)
181 
182    expect(models["gpt-5.4"]?.limit).toEqual(limit)
183    expect(models["gpt-5.5"]?.limit).toEqual({ context: 400_000, input: 272_000, output: 128_000 })
184    expect(models["gpt-5.6-sol"]?.limit).toEqual({ context: 500_000, input: 372_000, output: 128_000 })
185    expect(models["gpt-5.6-terra"]?.limit).toEqual({ context: 500_000, input: 372_000, output: 128_000 })
186    expect(models["gpt-5.6-luna"]?.limit).toEqual({ context: 500_000, input: 372_000, output: 128_000 })
187    expect(models["gpt-5.4-pro"]).toBeUndefined()
188    expect(models["gpt-5.7-pro"]).toBeDefined()
189    expect(models["gpt-5.6-sol-high"]).toBeDefined()
190    expect(await hooks.provider!.models!(provider as never, { auth: { type: "api" } } as never)).toBe(
191      provider.models as never,
192    )
193  })
194 
195  test("deduplicates concurrent Codex token refreshes", async () => {
196    let auth = {
197      type: "oauth" as const,
198      refresh: "refresh-old",
199      access: "",
200      expires: 0,
201    }
202    const authUpdates: Array<{
203      body: { refresh: string; access: string; expires: number; accountId?: string }
204    }> = []
205    let resolveRefresh: (() => void) | undefined
206    const refreshReady = new Promise<void>((resolve) => {
207      resolveRefresh = resolve
208    })
209    let refreshRequests = 0
210    const apiRequests: { authorization: string | null; accountId: string | null }[] = []
211 
212    using server = Bun.serve({
213      port: 0,
214      async fetch(request) {
215        const url = new URL(request.url)
216        if (url.pathname === "/oauth/token") {
217          expect(await request.text()).toContain("refresh_token=refresh-old")
218          refreshRequests += 1
219          await refreshReady
220          return Response.json({
221            id_token: createTestJwt({ chatgpt_account_id: "acc-123" }),
222            access_token: "access-new",
223            refresh_token: "refresh-new",
224            expires_in: 3600,
225          })
226        }
227 
228        if (url.pathname === "/backend-api/codex/responses") {
229          apiRequests.push({
230            authorization: request.headers.get("authorization"),
231            accountId: request.headers.get("ChatGPT-Account-Id"),
232          })
233          return new Response("{}", { status: 200 })
234        }
235 
236        return new Response("unexpected request", { status: 500 })
237      },
238    })
239 
240    const hooks = await CodexAuthPlugin(
241      {
242        client: {
243          auth: {
244            async set(input: { body: { refresh: string; access: string; expires: number; accountId?: string } }) {
245              authUpdates.push(input)
246              auth = {
247                type: "oauth",
248                refresh: input.body.refresh,
249                access: input.body.access,
250                expires: input.body.expires,
251                ...(input.body.accountId && { accountId: input.body.accountId }),
252              }
253            },
254          },
255        } as never,
256        project: {} as never,
257        directory: "",
258        worktree: "",
259        experimental_workspace: {
260          register() {},
261        },
262        serverUrl: new URL("https://example.com"),
263        $: {} as never,
264      },
265      {
266        issuer: server.url.origin,
267        codexApiEndpoint: new URL("/backend-api/codex/responses", server.url).toString(),
268      },
269    )
270    const loaded = await hooks.auth!.loader!(async () => auth as never, {} as never)
271 
272    const first = loaded.fetch!("https://api.openai.com/v1/responses")
273    const second = loaded.fetch!("https://api.openai.com/v1/responses")
274 
275    await waitFor(() => refreshRequests === 1)
276    expect(apiRequests).toHaveLength(0)
277 
278    resolveRefresh!()
279    await Promise.all([first, second])
280 
281    expect(refreshRequests).toBe(1)
282    expect(authUpdates).toHaveLength(1)
283    expect(authUpdates[0]?.body.refresh).toBe("refresh-new")
284    expect(authUpdates[0]?.body.access).toBe("access-new")
285    expect(authUpdates[0]?.body.accountId).toBe("acc-123")
286    expect(apiRequests).toEqual([
287      { authorization: "Bearer access-new", accountId: "acc-123" },
288      { authorization: "Bearer access-new", accountId: "acc-123" },
289    ])
290  })
291})
292 
293async function waitFor(predicate: () => boolean) {
294  const started = Date.now()
295  while (!predicate()) {
296    if (Date.now() - started > 1_000) throw new Error("timed out waiting for condition")
297    await new Promise((resolve) => setTimeout(resolve, 1))
298  }
299}
300