Fixing oauth flow on chat example (#27)

* wip

* Successfully getting a token from Groq!

* Simplify PKCE flow to use single key per provider

Previously, the PKCE flow stored multiple keys with timestamps (pkce_${providerId}_${timestamp})
and had complex logic to find the most recent key. This simplified implementation:

- Changes storage key format from `pkce_${providerId}_${timestamp}` to `pkce_${providerId}`
- Removes timestamp-based key sorting logic in completeOAuthFlow()
- Ensures only one PKCE key is stored per provider at any time
- Eliminates potential race conditions with multiple concurrent auth flows
- Simplifies debugging by having predictable key names

This applies to both OpenRouter and Groq OAuth flows. Also fixes TypeScript
build error by removing unused 'state' variable in OAuthCallback component.

Amp-Thread: https://ampcode.com/threads/T-12223468-a950-40a8-9d44-2e1a66e6efbe
Co-authored-by: Amp <amp@ampcode.com>

* Amp getting confused

* Able to run inference against staging Groq

* The old oauth callback got clobbered, try to restore it

* Restored, but still buggy

* updated data

* Pruning Groq models

* deploying

---------

Co-authored-by: Amp <amp@ampcode.com>

examples/chat-ui/src/App.tsx

@@ -1,13 +1,15 @@
 import { BrowserRouter as Router, Routes, Route } from 'react-router-dom'
 import ChatApp from './components/ChatApp'
-import OAuthCallback from './components/OAuthCallback'
+import PkceCallback from './components/PkceCallback.tsx'
+import { OAuthCallback } from './components/OAuthCallback.tsx'
 
 function App() {
   return (
     <Router>
       <Routes>
-        <Route path="/oauth/openrouter/callback" element={<OAuthCallback provider="openrouter" />} />
-        <Route path="/oauth/callback" element={<OAuthCallback provider="openrouter" />} />
+        <Route path="/oauth/groq/callback" element={<PkceCallback provider="groq" />} />
+        <Route path="/oauth/openrouter/callback" element={<PkceCallback provider="openrouter" />} />
+        <Route path="/oauth/callback" element={<OAuthCallback />} />
         <Route path="/" element={<ChatApp />} />
       </Routes>
     </Router>

examples/chat-ui/src/components/ChatApp.tsx

@@ -33,7 +33,9 @@ const ChatApp: React.FC<ChatAppProps> = () => {
   // Handle OAuth success messages from popups
   useEffect(() => {
     const handleMessage = (event: MessageEvent) => {
+      console.log('DEBUG: Received message in parent window:', event.data)
       if (event.data.type === 'oauth_success') {
+        console.log('DEBUG: OAuth success message received, triggering API key update')
         handleApiKeyUpdate()
       }
     }
@@ -42,6 +44,96 @@ const ChatApp: React.FC<ChatAppProps> = () => {
     return () => window.removeEventListener('message', handleMessage)
   }, [])
 
+  // Poll for OAuth token changes (fallback for when popup messaging doesn't work)
+  useEffect(() => {
+    const oauthProviders = ['groq', 'openrouter'] as const
+    let initialTokens: Record<string, string | null> = {}
+
+    // Capture initial token state
+    const captureInitialTokens = () => {
+      for (const providerId of oauthProviders) {
+        const tokenKey = `aiChatTemplate_token_${providerId}`
+        initialTokens[providerId] = localStorage.getItem(tokenKey)
+      }
+    }
+
+    const checkForNewTokens = () => {
+      for (const providerId of oauthProviders) {
+        const tokenKey = `aiChatTemplate_token_${providerId}`
+        const currentToken = localStorage.getItem(tokenKey)
+
+        // Check if token was added or changed
+        if (currentToken !== initialTokens[providerId]) {
+          try {
+            const parsedToken = JSON.parse(currentToken || '{}')
+            if (parsedToken.access_token) {
+              console.log('DEBUG: New OAuth token detected for', providerId, 'via polling')
+              handleApiKeyUpdate()
+              return true // Stop polling once we find a new token
+            }
+          } catch (e) {
+            // Invalid token format, continue checking
+          }
+        }
+      }
+      return false
+    }
+
+    let pollInterval: number | null = null
+
+    const startPolling = () => {
+      // Capture initial state
+      captureInitialTokens()
+
+      // Check immediately for existing valid tokens (in case we just redirected from OAuth)
+      console.log('DEBUG: Checking for existing OAuth tokens on startup')
+      for (const providerId of oauthProviders) {
+        const tokenKey = `aiChatTemplate_token_${providerId}`
+        const currentToken = localStorage.getItem(tokenKey)
+
+        console.log(`DEBUG: Checking ${providerId} token:`, currentToken ? 'exists' : 'not found')
+
+        if (currentToken) {
+          try {
+            const parsedToken = JSON.parse(currentToken)
+            console.log(`DEBUG: Parsed ${providerId} token:`, parsedToken)
+            if (parsedToken.access_token) {
+              console.log('DEBUG: Found existing OAuth token for', providerId, 'on startup')
+              handleApiKeyUpdate()
+              return // Don't start polling if we found a valid token
+            }
+          } catch (e) {
+            console.log(`DEBUG: Failed to parse ${providerId} token:`, e)
+          }
+        }
+      }
+
+      // Start polling every 500ms for new tokens
+      pollInterval = setInterval(() => {
+        if (checkForNewTokens()) {
+          if (pollInterval) {
+            clearInterval(pollInterval)
+            pollInterval = null
+            console.log('DEBUG: Stopped polling for OAuth tokens')
+          }
+        }
+      }, 500)
+
+      console.log('DEBUG: Started polling for OAuth token changes')
+    }
+
+    // Start polling after a short delay to allow for popup messages first
+    const delayedStart = setTimeout(startPolling, 100)
+
+    return () => {
+      if (delayedStart) clearTimeout(delayedStart)
+      if (pollInterval) {
+        clearInterval(pollInterval)
+        console.log('DEBUG: Cleaned up OAuth token polling')
+      }
+    }
+  }, [])
+
   const handleModelChange = (model: Model) => {
     setSelectedModel(model)
     saveSelectedModel(model)

examples/chat-ui/src/components/ModelSelector.tsx

@@ -39,7 +39,9 @@ const ModelSelector: React.FC<ModelSelectorProps> = ({ selectedModel, onModelCha
   // Handle OAuth success - show provider models when OAuth completes
   useEffect(() => {
     const handleMessage = (event: MessageEvent) => {
+      console.log('DEBUG: ModelSelector received message:', event.data)
       if (event.data.type === 'oauth_success' && event.data.provider) {
+        console.log('DEBUG: ModelSelector opening provider models modal for:', event.data.provider)
         const provider = providers[event.data.provider as keyof typeof providers]
         if (provider) {
           setProviderModelsModal({ isOpen: true, provider })

examples/chat-ui/src/components/OAuthCallback.tsx

@@ -1,124 +1,22 @@
-import React, { useEffect, useState, useRef } from 'react'
-import { useSearchParams } from 'react-router-dom'
-import { completeOAuthFlow } from '../utils/auth'
-import { SupportedProvider } from '../types/models'
-
-interface OAuthCallbackProps {
-  provider: SupportedProvider
-}
-
-const OAuthCallback: React.FC<OAuthCallbackProps> = ({ provider }) => {
-  const [searchParams] = useSearchParams()
-  const [status, setStatus] = useState<'loading' | 'success' | 'error'>('loading')
-  const [error, setError] = useState<string | null>(null)
-  const executedRef = useRef(false)
+import { useEffect } from 'react'
+import { onMcpAuthorization } from 'use-mcp'
 
+export function OAuthCallback() {
   useEffect(() => {
-    const handleCallback = async () => {
-      if (executedRef.current) {
-        console.log('DEBUG: Skipping duplicate OAuth callback execution')
-        return
-      }
-      executedRef.current = true
-
-      try {
-        const code = searchParams.get('code')
-        const state = searchParams.get('state')
-        const error = searchParams.get('error')
-
-        if (error) {
-          throw new Error(`OAuth error: ${error}`)
-        }
-
-        if (!code) {
-          throw new Error('Missing authorization code')
-        }
-
-        // OpenRouter doesn't use state parameter, but other providers might
-        const stateToUse = state || 'no-state'
-        await completeOAuthFlow(provider, code, stateToUse)
-        setStatus('success')
-
-        // Close popup after successful authentication
-        // Give extra time for debugging in development
-        setTimeout(() => {
-          if (window.opener) {
-            window.opener.postMessage({ type: 'oauth_success', provider }, '*')
-            window.close()
-          } else {
-            // Redirect to main page if not in popup
-            window.location.href = '/'
-          }
-        }, 3000)
-      } catch (err) {
-        console.error('OAuth callback error:', err)
-        setError(err instanceof Error ? err.message : 'Unknown error')
-        setStatus('error')
-      }
-    }
-
-    handleCallback()
-  }, [searchParams, provider])
+    onMcpAuthorization()
+  }, [])
 
   return (
-    <div className="min-h-screen flex items-center justify-center bg-gray-50">
-      <div className="max-w-md w-full bg-white shadow-lg rounded-lg p-8">
-        <div className="text-center">
-          {status === 'loading' && (
-            <>
-              <div className="animate-spin rounded-full h-12 w-12 border-b-2 border-blue-500 mx-auto mb-4"></div>
-              <h2 className="text-xl font-semibold text-gray-900 mb-2">Completing Authentication</h2>
-              <p className="text-gray-600">Connecting to {provider}...</p>
-            </>
-          )}
+    <div className="min-h-screen animated-bg-container flex items-center justify-center p-4">
+      <div className="bg-white rounded-lg shadow-xl p-8 text-center max-w-md w-full">
+        <h1 className="text-2xl font-bold text-zinc-900 mb-4">Authenticating...</h1>
+        <p className="text-zinc-600 mb-2">Please wait while we complete your authentication.</p>
+        <p className="text-sm text-zinc-500">This window should close automatically.</p>
 
-          {status === 'success' && (
-            <>
-              <div className="text-green-500 mb-4">
-                <svg className="w-12 h-12 mx-auto" fill="currentColor" viewBox="0 0 20 20">
-                  <path
-                    fillRule="evenodd"
-                    d="M10 18a8 8 0 100-16 8 8 0 000 16zm3.707-9.293a1 1 0 00-1.414-1.414L9 10.586 7.707 9.293a1 1 0 00-1.414 1.414l2 2a1 1 0 001.414 0l4-4z"
-                    clipRule="evenodd"
-                  />
-                </svg>
-              </div>
-              <h2 className="text-xl font-semibold text-gray-900 mb-2">Authentication Successful!</h2>
-              <p className="text-gray-600 mb-4">Successfully connected to {provider}. You can now close this window.</p>
-              <button
-                onClick={() => window.close()}
-                className="px-4 py-2 bg-blue-500 text-white rounded hover:bg-blue-600 transition-colors"
-              >
-                Close Window
-              </button>
-            </>
-          )}
-
-          {status === 'error' && (
-            <>
-              <div className="text-red-500 mb-4">
-                <svg className="w-12 h-12 mx-auto" fill="currentColor" viewBox="0 0 20 20">
-                  <path
-                    fillRule="evenodd"
-                    d="M10 18a8 8 0 100-16 8 8 0 000 16zM8.707 7.293a1 1 0 00-1.414 1.414L8.586 10l-1.293 1.293a1 1 0 101.414 1.414L10 11.414l1.293 1.293a1 1 0 001.414-1.414L11.414 10l1.293-1.293a1 1 0 00-1.414-1.414L10 8.586 8.707 7.293z"
-                    clipRule="evenodd"
-                  />
-                </svg>
-              </div>
-              <h2 className="text-xl font-semibold text-gray-900 mb-2">Authentication Failed</h2>
-              <p className="text-gray-600 mb-4">{error || 'An error occurred during authentication'}</p>
-              <button
-                onClick={() => window.close()}
-                className="px-4 py-2 bg-blue-500 text-white rounded hover:bg-blue-600 transition-colors"
-              >
-                Close Window
-              </button>
-            </>
-          )}
+        <div className="mt-6">
+          <div className="animate-spin rounded-full h-8 w-8 border-b-2 border-blue-600 mx-auto"></div>
         </div>
       </div>
     </div>
   )
 }
-
-export default OAuthCallback

examples/chat-ui/src/components/PkceCallback.tsx

@@ -0,0 +1,153 @@
+import React, { useEffect, useState, useRef } from 'react'
+import { useSearchParams } from 'react-router-dom'
+import { completeOAuthFlow } from '../utils/auth'
+import { SupportedProvider } from '../types/models'
+
+interface OAuthCallbackProps {
+  provider: SupportedProvider
+}
+
+const PkceCallback: React.FC<OAuthCallbackProps> = ({ provider }) => {
+  const [searchParams] = useSearchParams()
+  const [status, setStatus] = useState<'loading' | 'success' | 'error'>('loading')
+  const [error, setError] = useState<string | null>(null)
+  const executedRef = useRef(false)
+
+  useEffect(() => {
+    const handleCallback = async () => {
+      if (executedRef.current) {
+        console.log('DEBUG: Skipping duplicate OAuth callback execution')
+        return
+      }
+      executedRef.current = true
+
+      try {
+        const code = searchParams.get('code')
+        const error = searchParams.get('error')
+
+        if (error) {
+          throw new Error(`OAuth error: ${error}`)
+        }
+
+        if (!code) {
+          throw new Error('Missing authorization code')
+        }
+
+        // TODO: Add state parameter handling back if needed later
+        // const stateToUse = state || 'no-state'
+        await completeOAuthFlow(provider, code)
+        setStatus('success')
+
+        console.log('DEBUG: OAuth flow completed successfully')
+        console.log('DEBUG: window.opener exists:', !!window.opener)
+        console.log('DEBUG: window.opener closed:', window.opener?.closed)
+        console.log('DEBUG: window.parent exists:', !!window.parent)
+        console.log('DEBUG: window.parent === window:', window.parent === window)
+
+        // Try to send success message to parent if possible
+        const sendSuccessMessage = () => {
+          const message = { type: 'oauth_success', provider }
+
+          // Try window.opener first
+          if (window.opener && !window.opener.closed) {
+            console.log('DEBUG: Sending message via window.opener')
+            window.opener.postMessage(message, '*')
+            return true
+          }
+
+          // Also try window.parent as fallback
+          if (window.parent && window.parent !== window) {
+            console.log('DEBUG: Sending message via window.parent')
+            window.parent.postMessage(message, '*')
+            return true
+          }
+
+          return false
+        }
+
+        // Try to send success message
+        const messageSent = sendSuccessMessage()
+        console.log({ messageSent })
+
+        console.log('DEBUG: Closing popup in 100ms')
+        setTimeout(() => {
+          console.log('DEBUG: Attempting to close popup')
+          window.close()
+        }, 100)
+      } catch (err) {
+        console.error('OAuth callback error:', err)
+        setError(err instanceof Error ? err.message : 'Unknown error')
+        setStatus('error')
+      }
+    }
+
+    handleCallback()
+  }, [searchParams, provider])
+
+  return (
+    <div className="min-h-screen flex items-center justify-center bg-gray-50">
+      <div className="max-w-md w-full bg-white shadow-lg rounded-lg p-8">
+        <div className="text-center">
+          {status === 'loading' && (
+            <>
+              <div className="animate-spin rounded-full h-12 w-12 border-b-2 border-blue-500 mx-auto mb-4"></div>
+              <h2 className="text-xl font-semibold text-gray-900 mb-2">Completing Authentication</h2>
+              <p className="text-gray-600">Connecting to {provider}...</p>
+            </>
+          )}
+
+          {status === 'success' && (
+            <>
+              <div className="text-green-500 mb-4">
+                <svg className="w-12 h-12 mx-auto" fill="currentColor" viewBox="0 0 20 20">
+                  <path
+                    fillRule="evenodd"
+                    d="M10 18a8 8 0 100-16 8 8 0 000 16zm3.707-9.293a1 1 0 00-1.414-1.414L9 10.586 7.707 9.293a1 1 0 00-1.414 1.414l2 2a1 1 0 001.414 0l4-4z"
+                    clipRule="evenodd"
+                  />
+                </svg>
+              </div>
+              <h2 className="text-xl font-semibold text-gray-900 mb-2">Authentication Successful!</h2>
+              <p className="text-gray-600 mb-4">Successfully connected to {provider}.</p>
+              <div className="flex gap-2">
+                <button
+                  onClick={() => {
+                    window.opener?.postMessage({ type: 'oauth_success', provider }, '*')
+                    window.close()
+                  }}
+                  className="px-4 py-2 bg-blue-500 text-white rounded hover:bg-blue-600 transition-colors"
+                >
+                  Close Window
+                </button>
+              </div>
+            </>
+          )}
+
+          {status === 'error' && (
+            <>
+              <div className="text-red-500 mb-4">
+                <svg className="w-12 h-12 mx-auto" fill="currentColor" viewBox="0 0 20 20">
+                  <path
+                    fillRule="evenodd"
+                    d="M10 18a8 8 0 100-16 8 8 0 000 16zM8.707 7.293a1 1 0 00-1.414 1.414L8.586 10l-1.293 1.293a1 1 0 101.414 1.414L10 11.414l1.293 1.293a1 1 0 001.414-1.414L11.414 10l1.293-1.293a1 1 0 00-1.414-1.414L10 8.586 8.707 7.293z"
+                    clipRule="evenodd"
+                  />
+                </svg>
+              </div>
+              <h2 className="text-xl font-semibold text-gray-900 mb-2">Authentication Failed</h2>
+              <p className="text-gray-600 mb-4">{error || 'An error occurred during authentication'}</p>
+              <button
+                onClick={() => window.close()}
+                className="px-4 py-2 bg-blue-500 text-white rounded hover:bg-blue-600 transition-colors"
+              >
+                Close Window
+              </button>
+            </>
+          )}
+        </div>
+      </div>
+    </div>
+  )
+}
+
+export default PkceCallback

examples/chat-ui/src/hooks/useStreamResponse.ts

@@ -90,7 +90,7 @@ export const useStreamResponse = ({
     switch (model.provider.id) {
       case 'groq': {
         const apiKey = authHeaders.Authorization?.replace('Bearer ', '')
-        const groqProvider = createGroq({ apiKey })
+        const groqProvider = createGroq({ apiKey, baseURL: model.provider.baseUrl })
         baseModel = groqProvider(model.modelId)
         break
       }

examples/chat-ui/src/types/models.ts

@@ -67,13 +67,22 @@ export const providers: Record<SupportedProvider, Provider> = {
     baseUrl: 'https://api.groq.com/openai/v1',
     logo: '🚀',
     documentationUrl: 'https://console.groq.com/docs',
-    authType: 'apiKey',
     apiKeyHeader: 'Authorization',
-    // oauth: {
-    //   authorizeUrl: 'http://localhost:3000/keys/request',
-    //   tokenUrl: 'https://openrouter.ai/api/v1/auth/keys'
-    // },
+    authType: 'apiKey',
   },
+  // groq: {
+  //   id: 'groq',
+  //   name: 'Groq',
+  //   baseUrl: 'http://localhost:8000/api/openai/v1',
+  //   logo: '🚀',
+  //   documentationUrl: 'https://console.groq.com/docs',
+  //   apiKeyHeader: 'Authorization',
+  //   authType: 'oauth',
+  //   oauth: {
+  //     authorizeUrl: 'http://localhost:3000/keys/request',
+  //     tokenUrl: 'http://localhost:3000/keys/request/exchange',
+  //   },
+  // },
   anthropic: {
     id: 'anthropic',
     name: 'Anthropic',

examples/chat-ui/src/utils/auth.ts

@@ -11,7 +11,8 @@ export interface OAuthToken {
 // Types for PKCE flow
 interface PKCEState {
   code_verifier: string
-  state: string
+  // TODO: Add state support back if needed later
+  // state: string
 }
 
 // Generate a random code verifier for PKCE
@@ -35,15 +36,15 @@ async function generateCodeChallenge(verifier: string): Promise<string> {
     .replace(/=/g, '')
 }
 
-// Generate random state parameter
-function generateState(): string {
-  const array = new Uint8Array(16)
-  crypto.getRandomValues(array)
-  return btoa(String.fromCharCode(...array))
-    .replace(/\+/g, '-')
-    .replace(/\//g, '_')
-    .replace(/=/g, '')
-}
+// TODO: Add state generation back if needed later
+// function generateState(): string {
+//   const array = new Uint8Array(16)
+//   crypto.getRandomValues(array)
+//   return btoa(String.fromCharCode(...array))
+//     .replace(/\+/g, '-')
+//     .replace(/\//g, '_')
+//     .replace(/=/g, '')
+// }
 
 // API Key functions (existing functionality)
 export function hasApiKey(providerId: SupportedProvider): boolean {
@@ -111,11 +112,17 @@ export async function beginOAuthFlow(providerId: SupportedProvider): Promise<voi
 
   const codeVerifier = generateCodeVerifier()
   const codeChallenge = await generateCodeChallenge(codeVerifier)
-  const state = generateState()
 
-  // Store PKCE state in sessionStorage
-  const pkceState: PKCEState = { code_verifier: codeVerifier, state }
-  const storageKey = providerId === 'openrouter' ? `pkce_${providerId}_${Date.now()}` : `pkce_${providerId}_${state}`
+  console.log('DEBUG: Generated PKCE values for', providerId, {
+    codeVerifier: codeVerifier,
+    codeChallenge: codeChallenge,
+    verifierLength: codeVerifier.length,
+    challengeLength: codeChallenge.length,
+  })
+
+  // Store PKCE state in sessionStorage (one key per provider)
+  const pkceState: PKCEState = { code_verifier: codeVerifier }
+  const storageKey = `pkce_${providerId}`
   sessionStorage.setItem(storageKey, JSON.stringify(pkceState))
 
   // Construct authorization URL based on provider
@@ -132,60 +139,59 @@ export async function beginOAuthFlow(providerId: SupportedProvider): Promise<voi
   }
   authUrl.searchParams.set('code_challenge', codeChallenge)
   authUrl.searchParams.set('code_challenge_method', 'S256')
+  // TODO: Add state parameter back if needed later
+  // authUrl.searchParams.set('state', state)
 
   // Open popup or redirect
-  const popup = window.open(authUrl.toString(), `oauth_${providerId}`, 'width=600,height=700')
+  console.log('DEBUG: Opening OAuth popup for', providerId, 'with URL:', authUrl.toString())
+  const popup = window.open(authUrl.toString(), `oauth_${providerId}`, 'width=600,height=700,popup=1')
 
   if (!popup) {
-    throw new Error('Failed to open OAuth popup. Please allow popups for this site.')
+    console.log('DEBUG: Failed to open popup, trying redirect in same window')
+    // Fallback to redirect in same window
+    window.location.href = authUrl.toString()
+    return
   }
+
+  console.log('DEBUG: Popup opened successfully:', popup)
+  console.log('DEBUG: Popup closed:', popup.closed)
+  console.log('DEBUG: Popup window object:', popup.window)
 }
 
-export async function completeOAuthFlow(providerId: SupportedProvider, code: string, state: string): Promise<void> {
-  console.log('DEBUG: Starting OAuth completion for', providerId, 'with code:', code?.substring(0, 10) + '...', 'state:', state)
+export async function completeOAuthFlow(providerId: SupportedProvider, code: string): Promise<void> {
+  console.log('DEBUG: Starting OAuth completion for', providerId, 'with code:', code?.substring(0, 10) + '...')
+  console.log('DEBUG: Full code for debugging:', code)
+  console.log('DEBUG: Provider config:', providers[providerId])
 
   const provider = providers[providerId]
   if (!provider.oauth) {
     throw new Error(`Provider ${providerId} does not support OAuth`)
   }
 
-  // Retrieve PKCE state
-  let pkceState: PKCEState
+  // Retrieve PKCE state (single key per provider)
+  const storageKey = `pkce_${providerId}`
+  const pkceStateJson = sessionStorage.getItem(storageKey)
 
-  if (state === 'no-state' && providerId === 'openrouter') {
-    // OpenRouter doesn't use state, find the most recent PKCE state for this provider
-    const allKeys = Object.keys(sessionStorage)
-    const pkceKeys = allKeys.filter((key) => key.startsWith(`pkce_${providerId}_`))
+  console.log('DEBUG: Looking for PKCE key:', storageKey)
 
-    console.log('DEBUG: Found PKCE keys:', pkceKeys)
-
-    if (pkceKeys.length === 0) {
-      throw new Error('PKCE state not found. Please try again.')
-    }
-
-    // Use the most recent one (sort by timestamp)
-    const sortedKeys = pkceKeys.sort((a, b) => {
-      const aTime = parseInt(a.split('_').pop() || '0')
-      const bTime = parseInt(b.split('_').pop() || '0')
-      return bTime - aTime // Most recent first
-    })
-
-    const pkceStateJson = sessionStorage.getItem(sortedKeys[0])!
-    pkceState = JSON.parse(pkceStateJson)
-
-    console.log('DEBUG: Using PKCE state:', { key: sortedKeys[0], state: pkceState })
-
-    // Clean up the state
-    sessionStorage.removeItem(sortedKeys[0])
-  } else {
-    const pkceStateJson = sessionStorage.getItem(`pkce_${providerId}_${state}`)
-    if (!pkceStateJson) {
-      throw new Error('PKCE state not found. Please try again.')
-    }
-    pkceState = JSON.parse(pkceStateJson)
-    console.log('DEBUG: Using PKCE state for state', state, ':', pkceState)
+  if (!pkceStateJson) {
+    throw new Error('PKCE state not found. Please try again.')
   }
 
+  const pkceState: PKCEState = JSON.parse(pkceStateJson)
+
+  console.log('DEBUG: Using PKCE state:', { key: storageKey, state: pkceState })
+  console.log('DEBUG: Code verifier length:', pkceState.code_verifier.length)
+  console.log('DEBUG: Code verifier sample:', pkceState.code_verifier.substring(0, 20) + '...')
+
+  // Test: regenerate challenge from verifier to verify it matches server expectation
+  const recomputedChallenge = await generateCodeChallenge(pkceState.code_verifier)
+  console.log('DEBUG: Recomputed challenge from verifier:', recomputedChallenge)
+  console.log('DEBUG: Server reported challenge was: dW2iEvNljlkhcRcryo3Z0GITcJM1liKcHlB5v8CDEu8')
+
+  // Clean up the state
+  sessionStorage.removeItem(storageKey)
+
   // Exchange code for token
   let tokenResponse: Response
 
@@ -218,18 +224,23 @@ export async function completeOAuthFlow(providerId: SupportedProvider, code: str
       duration: `${endTime - startTime}ms`,
     })
   } else {
-    // Standard OAuth2 flow for other providers
+    // Standard OAuth2 flow for other providers (Groq)
     const requestBody = new URLSearchParams({
       grant_type: 'authorization_code',
       code,
       redirect_uri: getRedirectUri(providerId),
       code_verifier: pkceState.code_verifier,
     })
-    console.log('DEBUG: Standard OAuth token request:', {
+
+    console.log('DEBUG: Groq token request:', {
       url: provider.oauth.tokenUrl,
       body: Object.fromEntries(requestBody.entries()),
+      codeVerifierLength: pkceState.code_verifier.length,
+      codeVerifierSample: pkceState.code_verifier.substring(0, 20) + '...',
+      fullCodeVerifier: pkceState.code_verifier, // For debugging
     })
 
+    const startTime = performance.now()
     tokenResponse = await fetch(provider.oauth.tokenUrl, {
       method: 'POST',
       headers: {
@@ -237,11 +248,13 @@ export async function completeOAuthFlow(providerId: SupportedProvider, code: str
       },
       body: requestBody,
     })
+    const endTime = performance.now()
 
-    console.log('DEBUG: Standard OAuth token response:', {
+    console.log('DEBUG: Groq token response:', {
       status: tokenResponse.status,
       statusText: tokenResponse.statusText,
       headers: Object.fromEntries(tokenResponse.headers.entries()),
+      duration: `${endTime - startTime}ms`,
     })
   }
 
@@ -263,6 +276,12 @@ export async function completeOAuthFlow(providerId: SupportedProvider, code: str
       access_token: tokenData.key,
       token_type: 'Bearer',
     }
+  } else if (providerId === 'groq') {
+    // Groq returns { api_key: "..." }
+    token = {
+      access_token: tokenData.api_key,
+      token_type: 'Bearer',
+    }
   } else {
     // Standard OAuth2 response
     token = {
@@ -273,12 +292,10 @@ export async function completeOAuthFlow(providerId: SupportedProvider, code: str
     }
   }
 
+  console.log('DEBUG: Saving token for', providerId, ':', token)
   setOAuthToken(providerId, token)
 
-  // Clean up PKCE state (already cleaned up above for OpenRouter)
-  if (state !== 'no-state') {
-    sessionStorage.removeItem(`pkce_${providerId}_${state}`)
-  }
+  // PKCE state already cleaned up above
 }
 
 // Get authentication headers for API calls
@@ -318,3 +335,30 @@ function getRedirectUri(providerId: SupportedProvider): string {
   const baseUrl = window.location.origin
   return `${baseUrl}/oauth/${providerId}/callback`
 }
+
+// Test function to verify PKCE implementation with known values
+export async function testPKCEImplementation(): Promise<void> {
+  console.log('=== TESTING PKCE IMPLEMENTATION ===')
+
+  // Test with a known code verifier (from RFC 7636 example)
+  const testVerifier = 'dBjftJeZ4CVP-mB92K27uhbUJU1p1r_wW1gFWFOEjXk'
+  const expectedChallenge = 'E9Melhoa2OwvFrEMTJguCHaoeK1t8URWbuGJSstw-cM'
+
+  const computedChallenge = await generateCodeChallenge(testVerifier)
+
+  console.log('DEBUG: Test verifier:', testVerifier)
+  console.log('DEBUG: Expected challenge:', expectedChallenge)
+  console.log('DEBUG: Computed challenge:', computedChallenge)
+  console.log('DEBUG: Challenges match:', computedChallenge === expectedChallenge)
+
+  // Test with current implementation
+  const currentVerifier = generateCodeVerifier()
+  const currentChallenge = await generateCodeChallenge(currentVerifier)
+
+  console.log('DEBUG: Current verifier:', currentVerifier)
+  console.log('DEBUG: Current challenge:', currentChallenge)
+  console.log('DEBUG: Verifier length:', currentVerifier.length)
+  console.log('DEBUG: Challenge length:', currentChallenge.length)
+
+  console.log('=== END PKCE TEST ===')
+}