Add Codex OAuth provider flow for ChatGPT account sign-in (#503)

* feat: add Codex OAuth provider flow

* fix: harden Codex OAuth storage, session activation, and UI

src/services/oauth/auth-code-listener.analytics.test.ts

@@ -0,0 +1,155 @@
+import { afterEach, expect, mock, test } from 'bun:test'
+
+afterEach(() => {
+  mock.restore()
+})
+
+test('custom error responses log the error redirect analytics event', async () => {
+  const events: Array<{
+    name: string
+    metadata: Record<string, boolean | number | undefined>
+  }> = []
+
+  mock.module('src/services/analytics/index.js', () => ({
+    logEvent: (
+      name: string,
+      metadata: Record<string, boolean | number | undefined>,
+    ) => {
+      events.push({ name, metadata })
+    },
+  }))
+
+  const { AuthCodeListener } = await import(
+    `./auth-code-listener.js?ts=${Date.now()}-${Math.random()}`
+  )
+  const listener = new AuthCodeListener('/callback')
+  const response = {
+    writeHead: () => {},
+    end: () => {},
+  }
+
+  ;(listener as any).pendingResponse = response
+
+  listener.handleErrorRedirect(res => {
+    res.writeHead(400, {
+      'Content-Type': 'text/plain; charset=utf-8',
+    })
+    res.end('cancelled')
+  })
+
+  expect(events).toEqual([
+    {
+      name: 'tengu_oauth_automatic_redirect_error',
+      metadata: { custom_handler: true },
+    },
+  ])
+})
+
+test('custom handlers that do not end the response are closed automatically and still log analytics', async () => {
+  const events: Array<{
+    name: string
+    metadata: Record<string, boolean | number | undefined>
+  }> = []
+  const response = {
+    destroyed: false,
+    headersSent: false,
+    writableEnded: false,
+    writeHead: () => {
+      response.headersSent = true
+    },
+    end: () => {
+      response.writableEnded = true
+    },
+  }
+
+  mock.module('src/services/analytics/index.js', () => ({
+    logEvent: (
+      name: string,
+      metadata: Record<string, boolean | number | undefined>,
+    ) => {
+      events.push({ name, metadata })
+    },
+  }))
+
+  mock.module('../../utils/log.js', () => ({
+    logError: () => {},
+  }))
+
+  const { AuthCodeListener } = await import(
+    `./auth-code-listener.js?ts=${Date.now()}-${Math.random()}`
+  )
+  const listener = new AuthCodeListener('/callback')
+
+  ;(listener as any).pendingResponse = response
+
+  listener.handleErrorRedirect(res => {
+    res.writeHead(400, {
+      'Content-Type': 'text/plain; charset=utf-8',
+    })
+  })
+
+  expect(response.writableEnded).toBe(true)
+  expect((listener as any).pendingResponse).toBeNull()
+  expect(events).toEqual([
+    {
+      name: 'tengu_oauth_automatic_redirect_error',
+      metadata: { custom_handler: true },
+    },
+  ])
+})
+
+test('custom handlers that throw are logged, converted to a fallback response, and do not log analytics', async () => {
+  const events: Array<{
+    name: string
+    metadata: Record<string, boolean | number | undefined>
+  }> = []
+  const loggedErrors: unknown[] = []
+  const response = {
+    destroyed: false,
+    headersSent: false,
+    writableEnded: false,
+    statusCode: 0,
+    body: '',
+    writeHead: (statusCode: number) => {
+      response.headersSent = true
+      response.statusCode = statusCode
+    },
+    end: (body = '') => {
+      response.writableEnded = true
+      response.body = body
+    },
+  }
+
+  mock.module('src/services/analytics/index.js', () => ({
+    logEvent: (
+      name: string,
+      metadata: Record<string, boolean | number | undefined>,
+    ) => {
+      events.push({ name, metadata })
+    },
+  }))
+
+  mock.module('../../utils/log.js', () => ({
+    logError: (error: unknown) => {
+      loggedErrors.push(error)
+    },
+  }))
+
+  const { AuthCodeListener } = await import(
+    `./auth-code-listener.js?ts=${Date.now()}-${Math.random()}`
+  )
+  const listener = new AuthCodeListener('/callback')
+
+  ;(listener as any).pendingResponse = response
+
+  listener.handleErrorRedirect(() => {
+    throw new Error('handler exploded')
+  })
+
+  expect(response.statusCode).toBe(500)
+  expect(response.body).toBe('Authentication redirect failed')
+  expect(response.writableEnded).toBe(true)
+  expect((listener as any).pendingResponse).toBeNull()
+  expect(loggedErrors).toHaveLength(1)
+  expect(events).toEqual([])
+})

src/services/oauth/auth-code-listener.test.ts

@@ -0,0 +1,31 @@
+import { afterEach, expect, test } from 'bun:test'
+
+import { AuthCodeListener } from './auth-code-listener.js'
+
+const listeners: AuthCodeListener[] = []
+
+afterEach(() => {
+  while (listeners.length > 0) {
+    listeners.pop()?.close()
+  }
+})
+
+test('cancelPendingAuthorization rejects a pending OAuth wait', async () => {
+  const listener = new AuthCodeListener('/callback')
+  listeners.push(listener)
+
+  await listener.start()
+
+  const pendingAuthorization = listener.waitForAuthorization(
+    'state-test',
+    async () => {},
+  )
+
+  listener.cancelPendingAuthorization(
+    new Error('Codex OAuth flow was cancelled.'),
+  )
+
+  await expect(pendingAuthorization).rejects.toThrow(
+    'Codex OAuth flow was cancelled.',
+  )
+})

src/services/oauth/auth-code-listener.ts

@@ -71,6 +71,42 @@ export class AuthCodeListener {
     })
   }
 
+  private respondToPendingRequest(options: {
+    handler: (res: ServerResponse) => void
+    analyticsEvent:
+      | 'tengu_oauth_automatic_redirect'
+      | 'tengu_oauth_automatic_redirect_error'
+    analyticsMetadata?: Record<string, boolean>
+  }): void {
+    if (!this.pendingResponse) return
+
+    const response = this.pendingResponse
+    try {
+      options.handler(response)
+
+      if (!response.writableEnded && !response.destroyed) {
+        response.end()
+      }
+
+      logEvent(options.analyticsEvent, options.analyticsMetadata ?? {})
+    } catch (error) {
+      logError(error)
+
+      if (!response.headersSent && !response.destroyed) {
+        response.writeHead(500, {
+          'Content-Type': 'text/plain; charset=utf-8',
+        })
+      }
+      if (!response.writableEnded && !response.destroyed) {
+        response.end('Authentication redirect failed')
+      }
+    } finally {
+      if (this.pendingResponse === response) {
+        this.pendingResponse = null
+      }
+    }
+  }
+
   /**
    * Completes the OAuth flow by redirecting the user's browser to a success page.
    * Different success pages are shown based on the granted scopes.
@@ -85,9 +121,13 @@ export class AuthCodeListener {
 
     // If custom handler provided, use it instead of default redirect
     if (customHandler) {
-      customHandler(this.pendingResponse, scopes)
-      this.pendingResponse = null
-      logEvent('tengu_oauth_automatic_redirect', { custom_handler: true })
+      this.respondToPendingRequest({
+        handler: res => {
+          customHandler(res, scopes)
+        },
+        analyticsEvent: 'tengu_oauth_automatic_redirect',
+        analyticsMetadata: { custom_handler: true },
+      })
       return
     }
 
@@ -97,29 +137,48 @@ export class AuthCodeListener {
       : getOauthConfig().CONSOLE_SUCCESS_URL
 
     // Send browser to success page
-    this.pendingResponse.writeHead(302, { Location: successUrl })
-    this.pendingResponse.end()
-    this.pendingResponse = null
-
-    logEvent('tengu_oauth_automatic_redirect', {})
+    this.respondToPendingRequest({
+      handler: res => {
+        res.writeHead(302, { Location: successUrl })
+        res.end()
+      },
+      analyticsEvent: 'tengu_oauth_automatic_redirect',
+    })
   }
 
   /**
    * Handles error case by sending a redirect to the appropriate success page with an error indicator,
    * ensuring the browser flow is completed properly.
    */
-  handleErrorRedirect(): void {
+  handleErrorRedirect(customHandler?: (res: ServerResponse) => void): void {
     if (!this.pendingResponse) return
 
+    if (customHandler) {
+      this.respondToPendingRequest({
+        handler: customHandler,
+        analyticsEvent: 'tengu_oauth_automatic_redirect_error',
+        analyticsMetadata: { custom_handler: true },
+      })
+      return
+    }
+
     // TODO: swap to a different url once we have an error page
     const errorUrl = getOauthConfig().CLAUDEAI_SUCCESS_URL
 
-    // Send browser to error page
-    this.pendingResponse.writeHead(302, { Location: errorUrl })
-    this.pendingResponse.end()
-    this.pendingResponse = null
+    this.respondToPendingRequest({
+      handler: res => {
+        res.writeHead(302, { Location: errorUrl })
+        res.end()
+      },
+      analyticsEvent: 'tengu_oauth_automatic_redirect_error',
+    })
+  }
 
-    logEvent('tengu_oauth_automatic_redirect_error', {})
+  cancelPendingAuthorization(
+    error: Error = new Error('OAuth authorization was cancelled.'),
+  ): void {
+    this.reject(error)
+    this.close()
   }
 
   private startLocalListener(onReady: () => Promise<void>): void {
@@ -176,8 +235,7 @@ export class AuthCodeListener {
 
   private handleError(err: Error): void {
     logError(err)
-    this.close()
-    this.reject(err)
+    this.cancelPendingAuthorization(err)
   }
 
   private resolve(authorizationCode: string): void {
@@ -185,6 +243,7 @@ export class AuthCodeListener {
       this.promiseResolver(authorizationCode)
       this.promiseResolver = null
       this.promiseRejecter = null
+      this.expectedState = null
     }
   }
 
@@ -193,6 +252,7 @@ export class AuthCodeListener {
       this.promiseRejecter(error)
       this.promiseResolver = null
       this.promiseRejecter = null
+      this.expectedState = null
     }
   }
 
@@ -207,5 +267,8 @@ export class AuthCodeListener {
       this.localServer.removeAllListeners()
       this.localServer.close()
     }
+
+    this.expectedState = null
+    this.port = 0
   }
 }