feat(oauth): added token refresh to drive file fetch

app/api/auth/oauth/credentials/route.ts

@@ -1,5 +1,5 @@
 import { NextRequest, NextResponse } from 'next/server'
-import { and, eq, like } from 'drizzle-orm'
+import { and, eq } from 'drizzle-orm'
 import { jwtDecode } from 'jwt-decode'
 import { getSession } from '@/lib/auth'
 import { parseProvider } from '@/lib/oauth'

app/api/auth/oauth/drive/file/route.ts

@@ -0,0 +1,115 @@
+import { NextRequest, NextResponse } from 'next/server'
+import { eq } from 'drizzle-orm'
+import { getSession } from '@/lib/auth'
+import { refreshOAuthToken } from '@/lib/oauth'
+import { db } from '@/db'
+import { account } from '@/db/schema'
+
+/**
+ * Get a single file from Google Drive by ID
+ */
+export async function GET(request: NextRequest) {
+  try {
+    // Get the session
+    const session = await getSession()
+
+    // Check if the user is authenticated
+    if (!session?.user?.id) {
+      return NextResponse.json({ error: 'User not authenticated' }, { status: 401 })
+    }
+
+    // Get the credential ID and file ID from the query params
+    const { searchParams } = new URL(request.url)
+    const credentialId = searchParams.get('credentialId')
+    const fileId = searchParams.get('fileId')
+
+    if (!credentialId) {
+      return NextResponse.json({ error: 'Credential ID is required' }, { status: 400 })
+    }
+
+    if (!fileId) {
+      return NextResponse.json({ error: 'File ID is required' }, { status: 400 })
+    }
+
+    // Get the credential from the database
+    const credentials = await db.select().from(account).where(eq(account.id, credentialId)).limit(1)
+
+    if (!credentials.length) {
+      return NextResponse.json({ error: 'Credential not found' }, { status: 404 })
+    }
+
+    const credential = credentials[0]
+
+    // Check if the credential belongs to the user
+    if (credential.userId !== session.user.id) {
+      return NextResponse.json({ error: 'Unauthorized' }, { status: 403 })
+    }
+
+    // Check if the access token is valid
+    if (!credential.accessToken) {
+      return NextResponse.json({ error: 'No access token available' }, { status: 400 })
+    }
+
+    // Function to fetch file with a given token
+    const fetchFileWithToken = async (token: string) => {
+      const response = await fetch(
+        `https://www.googleapis.com/drive/v3/files/${fileId}?fields=id,name,mimeType,iconLink,webViewLink,thumbnailLink,createdTime,modifiedTime,size,owners`,
+        {
+          headers: {
+            Authorization: `Bearer ${token}`,
+          },
+        }
+      )
+      return response
+    }
+
+    // First attempt with current token
+    let response = await fetchFileWithToken(credential.accessToken)
+
+    // If unauthorized, try to refresh the token
+    if (response.status === 401 && credential.refreshToken) {
+      console.log('Access token expired, attempting to refresh...')
+
+      try {
+        // Refresh the token using the centralized utility
+        const refreshedToken = await refreshOAuthToken(
+          credential.providerId,
+          credential.refreshToken
+        )
+
+        if (refreshedToken) {
+          // Update the token in the database
+          await db
+            .update(account)
+            .set({
+              accessToken: refreshedToken,
+              accessTokenExpiresAt: new Date(Date.now() + 3600 * 1000), // Default 1 hour expiry
+              updatedAt: new Date(),
+            })
+            .where(eq(account.id, credentialId))
+
+          // Retry the request with the new token
+          response = await fetchFileWithToken(refreshedToken)
+        }
+      } catch (refreshError) {
+        console.error('Error refreshing token:', refreshError)
+        return NextResponse.json({ error: 'Failed to refresh access token' }, { status: 401 })
+      }
+    }
+
+    // Handle response
+    if (!response.ok) {
+      const error = await response.json().catch(() => ({ error: { message: 'Unknown error' } }))
+      return NextResponse.json(
+        { error: error.error?.message || 'Failed to fetch file from Google Drive' },
+        { status: response.status }
+      )
+    }
+
+    const file = await response.json()
+    return NextResponse.json({ file }, { status: 200 })
+  } catch (error) {
+    console.error('Error fetching file from Google Drive:', error)
+    return NextResponse.json({ error: 'Internal server error' }, { status: 500 })
+  }
+}

app/api/auth/oauth/drive/files/route.ts

@@ -1,6 +1,7 @@
 import { NextRequest, NextResponse } from 'next/server'
 import { eq } from 'drizzle-orm'
 import { getSession } from '@/lib/auth'
+import { refreshOAuthToken } from '@/lib/oauth'
 import { db } from '@/db'
 import { account } from '@/db/schema'
 
@@ -69,18 +70,55 @@ export async function GET(request: NextRequest) {
       }
     }
 
-    // Fetch files from Google Drive
-    const response = await fetch(
-      `https://www.googleapis.com/drive/v3/files?${queryParams}&fields=files(id,name,mimeType,iconLink,webViewLink,thumbnailLink,createdTime,modifiedTime,size,owners)`,
-      {
-        headers: {
-          Authorization: `Bearer ${credential.accessToken}`,
-        },
+    // Function to fetch files with a given token
+    const fetchFilesWithToken = async (token: string) => {
+      const response = await fetch(
+        `https://www.googleapis.com/drive/v3/files?${queryParams}&fields=files(id,name,mimeType,iconLink,webViewLink,thumbnailLink,createdTime,modifiedTime,size,owners)`,
+        {
+          headers: {
+            Authorization: `Bearer ${token}`,
+          },
+        }
+      )
+      return response
+    }
+
+    // First attempt with current token
+    let response = await fetchFilesWithToken(credential.accessToken)
+
+    // If unauthorized, try to refresh the token
+    if (response.status === 401 && credential.refreshToken) {
+      console.log('Access token expired, attempting to refresh...')
+
+      try {
+        // Refresh the token using the centralized utility
+        const refreshedToken = await refreshOAuthToken(
+          credential.providerId,
+          credential.refreshToken
+        )
+
+        if (refreshedToken) {
+          // Update the token in the database
+          await db
+            .update(account)
+            .set({
+              accessToken: refreshedToken,
+              accessTokenExpiresAt: new Date(Date.now() + 3600 * 1000), // Default 1 hour expiry
+              updatedAt: new Date(),
+            })
+            .where(eq(account.id, credentialId))
+
+          // Retry the request with the new token
+          response = await fetchFilesWithToken(refreshedToken)
+        }
+      } catch (refreshError) {
+        console.error('Error refreshing token:', refreshError)
+        return NextResponse.json({ error: 'Failed to refresh access token' }, { status: 401 })
       }
-    )
+    }
 
     if (!response.ok) {
-      const error = await response.json()
+      const error = await response.json().catch(() => ({ error: { message: 'Unknown error' } }))
       return NextResponse.json(
         { error: error.error?.message || 'Failed to fetch files from Google Drive' },
         { status: response.status }

app/api/auth/oauth/token/route.ts

@@ -1,7 +1,7 @@
 import { NextRequest, NextResponse } from 'next/server'
 import { and, eq } from 'drizzle-orm'
 import { getSession } from '@/lib/auth'
-import { client } from '@/lib/auth-client'
+import { refreshOAuthToken } from '@/lib/oauth'
 import { db } from '@/db'
 import { account, workflow } from '@/db/schema'
 
@@ -69,85 +69,31 @@ export async function POST(request: NextRequest) {
 
     if (needsRefresh && credential.refreshToken) {
       try {
-        // Get the provider from the providerId (e.g., 'google-email' -> 'google')
-        const provider = credential.providerId.split('-')[0]
+        const refreshedToken = await refreshOAuthToken(
+          credential.providerId,
+          credential.refreshToken
+        )
 
-        // Determine the token endpoint based on the provider
-        let tokenEndpoint: string
-        let clientId: string | undefined
-        let clientSecret: string | undefined
-
-        switch (provider) {
-          case 'google':
-            tokenEndpoint = 'https://oauth2.googleapis.com/token'
-            clientId = process.env.GOOGLE_CLIENT_ID
-            clientSecret = process.env.GOOGLE_CLIENT_SECRET
-            break
-          case 'github':
-            tokenEndpoint = 'https://github.com/login/oauth/access_token'
-            clientId = process.env.GITHUB_CLIENT_ID
-            clientSecret = process.env.GITHUB_CLIENT_SECRET
-            break
-          case 'x':
-            tokenEndpoint = 'https://api.x.com/2/oauth2/token'
-            clientId = process.env.X_CLIENT_ID
-            clientSecret = process.env.X_CLIENT_SECRET
-            break
-          default:
-            throw new Error(`Unsupported provider: ${provider}`)
+        if (!refreshedToken) {
+          throw new Error('Failed to refresh token')
         }
 
-        if (!clientId || !clientSecret) {
-          throw new Error(`Missing client credentials for provider: ${provider}`)
-        }
-
-        // Refresh the token
-        const response = await fetch(tokenEndpoint, {
-          method: 'POST',
-          headers: {
-            'Content-Type': 'application/x-www-form-urlencoded',
-            ...(provider === 'github' && {
-              Accept: 'application/json',
-            }),
-          },
-          body: new URLSearchParams({
-            client_id: clientId,
-            client_secret: clientSecret,
-            grant_type: 'refresh_token',
-            refresh_token: credential.refreshToken,
-          }).toString(),
-        })
-
-        if (!response.ok) {
-          const errorText = await response.text()
-          console.error('Token refresh failed:', response.status, errorText)
-          throw new Error(`Failed to refresh token: ${response.status} ${errorText}`)
-        }
-
-        const data = await response.json()
-
-        // Update the credential in the database
         await db
           .update(account)
           .set({
-            accessToken: data.access_token,
-            accessTokenExpiresAt: data.expires_in
-              ? new Date(Date.now() + data.expires_in * 1000)
-              : null,
-            refreshToken: data.refresh_token || credential.refreshToken, // Some providers don't return a new refresh token
+            accessToken: refreshedToken,
+            accessTokenExpiresAt: new Date(Date.now() + 3600 * 1000), // Default 1 hour expiry
             updatedAt: new Date(),
           })
           .where(eq(account.id, credentialId))
 
-        // Return the new access token
-        return NextResponse.json({ accessToken: data.access_token }, { status: 200 })
+        return NextResponse.json({ accessToken: refreshedToken }, { status: 200 })
       } catch (error) {
         console.error('Error refreshing token:', error)
         return NextResponse.json({ error: 'Failed to refresh access token' }, { status: 500 })
       }
     }
 
-    // Return the current access token
     return NextResponse.json({ accessToken: credential.accessToken }, { status: 200 })
   } catch (error) {
     console.error('Error getting access token:', error)

app/w/[id]/components/workflow-block/components/sub-block/components/credential-selector.tsx

@@ -1,7 +1,7 @@
 'use client'
 
 import { useCallback, useEffect, useRef, useState } from 'react'
-import { Check, ChevronDown, ExternalLink, Key, RefreshCw } from 'lucide-react'
+import { Check, ChevronDown, ExternalLink, RefreshCw } from 'lucide-react'
 import { Button } from '@/components/ui/button'
 import {
   Command,
@@ -19,7 +19,6 @@ import {
   OAUTH_PROVIDERS,
   OAuthProvider,
   getProviderIdFromServiceId,
-  getServiceByProviderAndId,
   getServiceIdFromScopes,
   parseProvider,
 } from '@/lib/oauth'

components/ui/file-selector.tsx

@@ -70,6 +70,7 @@ export function FileSelector({
   const [selectedFileId, setSelectedFileId] = useState(value)
   const [selectedFile, setSelectedFile] = useState<FileInfo | null>(null)
   const [isLoading, setIsLoading] = useState(false)
+  const [isLoadingSelectedFile, setIsLoadingSelectedFile] = useState(false)
   const [showOAuthModal, setShowOAuthModal] = useState(false)
   const initialFetchRef = useRef(false)
 
@@ -122,6 +123,42 @@ export function FileSelector({
     }
   }, [provider, getProviderId, selectedCredentialId])
 
+  // Fetch a single file by ID when we have a selectedFileId but no metadata
+  const fetchFileById = useCallback(
+    async (fileId: string) => {
+      if (!selectedCredentialId || !fileId) return null
+
+      setIsLoadingSelectedFile(true)
+      try {
+        // Construct query parameters
+        const queryParams = new URLSearchParams({
+          credentialId: selectedCredentialId,
+          fileId: fileId,
+        })
+
+        const response = await fetch(`/api/auth/oauth/drive/file?${queryParams.toString()}`)
+
+        if (response.ok) {
+          const data = await response.json()
+          if (data.file) {
+            setSelectedFile(data.file)
+            onFileInfoChange?.(data.file)
+            return data.file
+          }
+        } else {
+          console.error('Error fetching file by ID:', await response.text())
+        }
+        return null
+      } catch (error) {
+        console.error('Error fetching file by ID:', error)
+        return null
+      } finally {
+        setIsLoadingSelectedFile(false)
+      }
+    },
+    [selectedCredentialId, onFileInfoChange]
+  )
+
   // Fetch files from Google Drive
   const fetchFiles = useCallback(
     async (searchQuery?: string) => {
@@ -155,9 +192,9 @@ export function FileSelector({
               setSelectedFile(fileInfo)
               onFileInfoChange?.(fileInfo)
             } else if (!searchQuery) {
-              // Only reset if this is not a search query
-              setSelectedFile(null)
-              onFileInfoChange?.(null)
+              // Only try to fetch by ID if this is not a search query
+              // and we couldn't find the file in the list
+              fetchFileById(selectedFileId)
             }
           }
         } else {
@@ -171,7 +208,7 @@ export function FileSelector({
         setIsLoading(false)
       }
     },
-    [selectedCredentialId, mimeTypeFilter, selectedFileId, onFileInfoChange]
+    [selectedCredentialId, mimeTypeFilter, selectedFileId, onFileInfoChange, fetchFileById]
   )
 
   // Fetch credentials on initial mount
@@ -199,9 +236,19 @@ export function FileSelector({
         const fileInfo = files.find((file) => file.id === value) || null
         setSelectedFile(fileInfo)
         onFileInfoChange?.(fileInfo)
+      } else if (value && selectedCredentialId) {
+        // If we have a value but no files loaded yet, try to fetch the file by ID
+        fetchFileById(value)
       }
     }
-  }, [value, files, onFileInfoChange])
+  }, [value, files, onFileInfoChange, selectedCredentialId, fetchFileById])
+
+  // Try to fetch the file by ID when credentials become available
+  useEffect(() => {
+    if (selectedCredentialId && selectedFileId && !selectedFile) {
+      fetchFileById(selectedFileId)
+    }
+  }, [selectedCredentialId, selectedFileId, selectedFile, fetchFileById])
 
   // Handle file selection
   const handleSelectFile = (file: FileInfo) => {
@@ -335,6 +382,11 @@ export function FileSelector({
                   {getFileIcon(selectedFile, 'sm')}
                   <span className="font-normal truncate">{selectedFile.name}</span>
                 </div>
+              ) : selectedFileId && (isLoadingSelectedFile || !selectedCredentialId) ? (
+                <div className="flex items-center gap-2">
+                  <RefreshCw className="h-4 w-4 animate-spin" />
+                  <span className="text-muted-foreground">Loading document...</span>
+                </div>
               ) : (
                 <div className="flex items-center gap-2">
                   {getProviderIcon(provider)}

lib/oauth.ts

@@ -265,3 +265,78 @@ export function parseProvider(provider: OAuthProvider): ProviderConfig {
     featureType: 'default',
   }
 }
+
+/**
+ * Refresh an OAuth token
+ * This is a server-side utility function to refresh OAuth tokens
+ * @param providerId The provider ID (e.g., 'google-drive')
+ * @param refreshToken The refresh token to use
+ * @returns The new access token, or null if refresh failed
+ */
+export async function refreshOAuthToken(
+  providerId: string,
+  refreshToken: string
+): Promise<string | null> {
+  try {
+    // Get the provider from the providerId (e.g., 'google-drive' -> 'google')
+    const provider = providerId.split('-')[0]
+
+    // Determine the token endpoint based on the provider
+    let tokenEndpoint: string
+    let clientId: string | undefined
+    let clientSecret: string | undefined
+
+    switch (provider) {
+      case 'google':
+        tokenEndpoint = 'https://oauth2.googleapis.com/token'
+        clientId = process.env.GOOGLE_CLIENT_ID
+        clientSecret = process.env.GOOGLE_CLIENT_SECRET
+        break
+      case 'github':
+        tokenEndpoint = 'https://github.com/login/oauth/access_token'
+        clientId = process.env.GITHUB_CLIENT_ID
+        clientSecret = process.env.GITHUB_CLIENT_SECRET
+        break
+      case 'x':
+        tokenEndpoint = 'https://api.x.com/2/oauth2/token'
+        clientId = process.env.X_CLIENT_ID
+        clientSecret = process.env.X_CLIENT_SECRET
+        break
+      default:
+        throw new Error(`Unsupported provider: ${provider}`)
+    }
+
+    if (!clientId || !clientSecret) {
+      throw new Error(`Missing client credentials for provider: ${provider}`)
+    }
+
+    // Refresh the token
+    const response = await fetch(tokenEndpoint, {
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/x-www-form-urlencoded',
+        ...(provider === 'github' && {
+          Accept: 'application/json',
+        }),
+      },
+      body: new URLSearchParams({
+        client_id: clientId,
+        client_secret: clientSecret,
+        grant_type: 'refresh_token',
+        refresh_token: refreshToken,
+      }).toString(),
+    })
+
+    if (!response.ok) {
+      const errorText = await response.text()
+      console.error('Token refresh failed:', response.status, errorText)
+      throw new Error(`Failed to refresh token: ${response.status} ${errorText}`)
+    }
+
+    const data = await response.json()
+    return data.access_token || null
+  } catch (error) {
+    console.error('Error refreshing token:', error)
+    return null
+  }
+}