1  import { logEvent } from 'src/services/analytics/index.js'
  2  import { openBrowser } from '../../utils/browser.js'
  3  import { AuthCodeListener } from './auth-code-listener.js'
  4  import * as client from './client.js'
  5  import * as crypto from './crypto.js'
  6  import type {
  7    OAuthProfileResponse,
  8    OAuthTokenExchangeResponse,
  9    OAuthTokens,
 10    RateLimitTier,
 11    SubscriptionType,
 12  } from './types.js'
 13  
 14  /**
 15   * OAuth service that handles the OAuth 2.0 authorization code flow with PKCE.
 16   *
 17   * Supports two ways to get authorization codes:
 18   * 1. Automatic: Opens browser, redirects to localhost where we capture the code
 19   * 2. Manual: User manually copies and pastes the code (used in non-browser environments)
 20   */
 21  export class OAuthService {
 22    private codeVerifier: string
 23    private authCodeListener: AuthCodeListener | null = null
 24    private port: number | null = null
 25    private manualAuthCodeResolver: ((authorizationCode: string) => void) | null =
 26      null
 27  
 28    constructor() {
 29      this.codeVerifier = crypto.generateCodeVerifier()
 30    }
 31  
 32    async startOAuthFlow(
 33      authURLHandler: (url: string, automaticUrl?: string) => Promise<void>,
 34      options?: {
 35        loginWithClaudeAi?: boolean
 36        inferenceOnly?: boolean
 37        expiresIn?: number
 38        orgUUID?: string
 39        loginHint?: string
 40        loginMethod?: string
 41        /**
 42         * Don't call openBrowser(). Caller takes both URLs via authURLHandler
 43         * and decides how/where to open them. Used by the SDK control protocol
 44         * (claude_authenticate) where the SDK client owns the user's display,
 45         * not this process.
 46         */
 47        skipBrowserOpen?: boolean
 48      },
 49    ): Promise<OAuthTokens> {
 50      // Create OAuth callback listener and start it
 51      this.authCodeListener = new AuthCodeListener()
 52      this.port = await this.authCodeListener.start()
 53  
 54      // Generate PKCE values and state
 55      const codeChallenge = crypto.generateCodeChallenge(this.codeVerifier)
 56      const state = crypto.generateState()
 57  
 58      // Build auth URLs for both automatic and manual flows
 59      const opts = {
 60        codeChallenge,
 61        state,
 62        port: this.port,
 63        loginWithClaudeAi: options?.loginWithClaudeAi,
 64        inferenceOnly: options?.inferenceOnly,
 65        orgUUID: options?.orgUUID,
 66        loginHint: options?.loginHint,
 67        loginMethod: options?.loginMethod,
 68      }
 69      const manualFlowUrl = client.buildAuthUrl({ ...opts, isManual: true })
 70      const automaticFlowUrl = client.buildAuthUrl({ ...opts, isManual: false })
 71  
 72      // Wait for either automatic or manual auth code
 73      const authorizationCode = await this.waitForAuthorizationCode(
 74        state,
 75        async () => {
 76          if (options?.skipBrowserOpen) {
 77            // Hand both URLs to the caller. The automatic one still works
 78            // if the caller opens it on the same host (localhost listener
 79            // is running); the manual one works from anywhere.
 80            await authURLHandler(manualFlowUrl, automaticFlowUrl)
 81          } else {
 82            await authURLHandler(manualFlowUrl) // Show manual option to user
 83            await openBrowser(automaticFlowUrl) // Try automatic flow
 84          }
 85        },
 86      )
 87  
 88      // Check if the automatic flow is still active (has a pending response)
 89      const isAutomaticFlow = this.authCodeListener?.hasPendingResponse() ?? false
 90      logEvent('tengu_oauth_auth_code_received', { automatic: isAutomaticFlow })
 91  
 92      try {
 93        // Exchange authorization code for tokens
 94        const tokenResponse = await client.exchangeCodeForTokens(
 95          authorizationCode,
 96          state,
 97          this.codeVerifier,
 98          this.port!,
 99          !isAutomaticFlow, // Pass isManual=true if it's NOT automatic flow
100          options?.expiresIn,
101        )
102  
103        // Fetch profile info (subscription type and rate limit tier) for the
104        // returned OAuthTokens. Logout and account storage are handled by the
105        // caller (installOAuthTokens in auth.ts).
106        const profileInfo = await client.fetchProfileInfo(
107          tokenResponse.access_token,
108        )
109  
110        // Handle success redirect for automatic flow
111        if (isAutomaticFlow) {
112          const scopes = client.parseScopes(tokenResponse.scope)
113          this.authCodeListener?.handleSuccessRedirect(scopes)
114        }
115  
116        return this.formatTokens(
117          tokenResponse,
118          profileInfo.subscriptionType,
119          profileInfo.rateLimitTier,
120          profileInfo.rawProfile,
121        )
122      } catch (error) {
123        // If we have a pending response, send an error redirect before closing
124        if (isAutomaticFlow) {
125          this.authCodeListener?.handleErrorRedirect()
126        }
127        throw error
128      } finally {
129        // Always cleanup
130        this.authCodeListener?.close()
131      }
132    }
133  
134    private async waitForAuthorizationCode(
135      state: string,
136      onReady: () => Promise<void>,
137    ): Promise<string> {
138      return new Promise((resolve, reject) => {
139        // Set up manual auth code resolver
140        this.manualAuthCodeResolver = resolve
141  
142        // Start automatic flow
143        this.authCodeListener
144          ?.waitForAuthorization(state, onReady)
145          .then(authorizationCode => {
146            this.manualAuthCodeResolver = null
147            resolve(authorizationCode)
148          })
149          .catch(error => {
150            this.manualAuthCodeResolver = null
151            reject(error)
152          })
153      })
154    }
155  
156    // Handle manual flow callback when user pastes the auth code
157    handleManualAuthCodeInput(params: {
158      authorizationCode: string
159      state: string
160    }): void {
161      if (this.manualAuthCodeResolver) {
162        this.manualAuthCodeResolver(params.authorizationCode)
163        this.manualAuthCodeResolver = null
164        // Close the auth code listener since manual input was used
165        this.authCodeListener?.close()
166      }
167    }
168  
169    private formatTokens(
170      response: OAuthTokenExchangeResponse,
171      subscriptionType: SubscriptionType | null,
172      rateLimitTier: RateLimitTier | null,
173      profile?: OAuthProfileResponse,
174    ): OAuthTokens {
175      return {
176        accessToken: response.access_token,
177        refreshToken: response.refresh_token,
178        expiresAt: Date.now() + response.expires_in * 1000,
179        scopes: client.parseScopes(response.scope),
180        subscriptionType,
181        rateLimitTier,
182        profile,
183        tokenAccount: response.account
184          ? {
185              uuid: response.account.uuid,
186              emailAddress: response.account.email_address,
187              organizationUuid: response.organization?.uuid,
188            }
189          : undefined,
190      }
191    }
192  
193    // Clean up any resources (like the local server)
194    cleanup(): void {
195      this.authCodeListener?.close()
196      this.manualAuthCodeResolver = null
197    }
198  }