feat: Implement Google Drive integration with PKCE OAuth flow

app/(auth)/auth.config.ts

@@ -10,4 +10,5 @@ export const authConfig = {
     // while this file is also used in non-Node.js environments
   ],
   callbacks: {},
+  secret: process.env.AUTH_SECRET,
 } satisfies NextAuthConfig;

app/(auth)/auth.ts

@@ -37,6 +37,9 @@ export const {
   signOut,
 } = NextAuth({
   ...authConfig,
+  session: {
+    strategy: 'jwt',
+  },
   providers: [
     Credentials({
       credentials: {},

app/api/google-drive/auth/route.ts

@@ -0,0 +1,99 @@
+import { NextRequest, NextResponse } from 'next/server';
+import { GoogleDriveService } from '@/lib/google-drive';
+
+// Log the environment variables (without sensitive data)
+console.log('Auth Route - Environment Check:');
+console.log('NEXT_PUBLIC_APP_URL:', process.env.NEXT_PUBLIC_APP_URL);
+console.log('GOOGLE_CLIENT_ID exists:', !!process.env.GOOGLE_CLIENT_ID);
+console.log('GOOGLE_CLIENT_SECRET exists:', !!process.env.GOOGLE_CLIENT_SECRET);
+console.log('Redirect URI:', `${process.env.NEXT_PUBLIC_APP_URL}/api/google-drive/auth`);
+
+if (!process.env.GOOGLE_CLIENT_ID || !process.env.GOOGLE_CLIENT_SECRET || !process.env.NEXT_PUBLIC_APP_URL) {
+  throw new Error('Missing required environment variables for Google Drive integration');
+}
+
+const driveService = new GoogleDriveService(
+  process.env.GOOGLE_CLIENT_ID,
+  process.env.GOOGLE_CLIENT_SECRET,
+  `${process.env.NEXT_PUBLIC_APP_URL}/api/google-drive/auth`
+);
+
+export async function GET(request: NextRequest) {
+  try {
+    const searchParams = request.nextUrl.searchParams;
+    const code = searchParams.get('code');
+    const codeVerifier = searchParams.get('code_verifier');
+
+    console.log('Auth Route - Request:', {
+      url: request.url,
+      hasCode: !!code,
+      hasCodeVerifier: !!codeVerifier,
+      searchParams: Object.fromEntries(searchParams.entries())
+    });
+
+    if (!code) {
+      // Redirect to Google OAuth
+      const { url: authUrl, codeVerifier } = await driveService.getAuthUrl();
+      console.log('Auth Route - Redirecting to Google consent screen:', authUrl);
+
+      // Return a 307 Temporary Redirect with the code verifier
+      const response = NextResponse.redirect(authUrl);
+      response.headers.set('Cache-Control', 'no-store');
+      response.cookies.set('code_verifier', codeVerifier, {
+        httpOnly: true,
+        secure: process.env.NODE_ENV === 'production',
+        sameSite: 'lax',
+        maxAge: 3600 // 1 hour
+      });
+      return response;
+    }
+
+    // Get the code verifier from the cookie
+    const storedCodeVerifier = request.cookies.get('code_verifier')?.value;
+    if (!storedCodeVerifier) {
+      throw new Error('No code verifier found in cookies');
+    }
+
+    // Exchange code for access token
+    const accessToken = await driveService.getAccessToken(code, storedCodeVerifier);
+    console.log('Auth Route - Successfully obtained access token');
+
+    // Create a new URL for the home page
+    const homeUrl = new URL('/', request.url);
+    console.log('Auth Route - Redirecting to home page:', homeUrl.toString());
+
+    // Create the response with the redirect
+    const response = NextResponse.redirect(homeUrl);
+
+    // Set the access token cookie
+    response.cookies.set('google_drive_token', accessToken, {
+      httpOnly: true,
+      secure: process.env.NODE_ENV === 'production',
+      sameSite: 'lax',
+      maxAge: 3600 // 1 hour
+    });
+
+    // Clear the code verifier cookie
+    response.cookies.delete('code_verifier');
+
+    // Set cache control headers
+    response.headers.set('Cache-Control', 'no-store, no-cache, must-revalidate');
+    response.headers.set('Pragma', 'no-cache');
+    response.headers.set('Expires', '0');
+
+    return response;
+  } catch (error: any) {
+    console.error('Error in Google Drive auth:', {
+      error: error.message,
+      response: error.response?.data,
+      status: error.response?.status
+    });
+
+    // Create a new URL for the error page
+    const errorUrl = new URL('/?error=auth_failed', request.url);
+    console.log('Auth Route - Redirecting to error page:', errorUrl.toString());
+
+    // Return a redirect to the error page
+    return NextResponse.redirect(errorUrl);
+  }
+} 

app/api/google-drive/files/route.ts

@@ -0,0 +1,30 @@
+import { NextRequest, NextResponse } from 'next/server';
+import { GoogleDriveService } from '@/lib/google-drive';
+
+const driveService = new GoogleDriveService(
+  process.env.GOOGLE_CLIENT_ID!,
+  process.env.GOOGLE_CLIENT_SECRET!,
+  `${process.env.NEXT_PUBLIC_APP_URL}/api/google-drive/auth`
+);
+
+export async function GET(request: NextRequest) {
+  const token = request.cookies.get('google_drive_token')?.value;
+
+  if (!token) {
+    return NextResponse.json(
+      { error: 'Not authenticated' },
+      { status: 401 }
+    );
+  }
+
+  try {
+    const files = await driveService.listFiles();
+    return NextResponse.json({ files });
+  } catch (error) {
+    console.error('Error listing Google Drive files:', error);
+    return NextResponse.json(
+      { error: 'Failed to list files' },
+      { status: 500 }
+    );
+  }
+} 

components/chat.tsx

@@ -20,6 +20,8 @@ import { useSearchParams } from 'next/navigation';
 import { useChatVisibility } from '@/hooks/use-chat-visibility';
 import { useAutoResume } from '@/hooks/use-auto-resume';
 import { ChatSDKError } from '@/lib/errors';
+import { GoogleDrivePicker } from './google-drive-picker';
+import { GoogleDriveFile } from '@/lib/google-drive';
 
 export function Chat({
   id,
@@ -107,6 +109,14 @@ export function Chat({
 
   const [attachments, setAttachments] = useState<Array<Attachment>>([]);
   const isArtifactVisible = useArtifactSelector((state) => state.isVisible);
+  const [selectedFile, setSelectedFile] = useState<GoogleDriveFile | null>(null);
+
+  const handleFileSelect = (file: GoogleDriveFile) => {
+    setSelectedFile(file);
+    // Add the file content to the chat context
+    const fileContent = `File: ${file.name}\nType: ${file.mimeType}\nLink: ${file.webViewLink}`;
+    // You can add this content to your chat context or send it as a message
+  };
 
   useAutoResume({
     autoResume,
@@ -140,20 +150,23 @@ export function Chat({
 
         <form className="flex mx-auto px-4 bg-background pb-4 md:pb-6 gap-2 w-full md:max-w-3xl">
           {!isReadonly && (
-            <MultimodalInput
-              chatId={id}
-              input={input}
-              setInput={setInput}
-              handleSubmit={handleSubmit}
-              status={status}
-              stop={stop}
-              attachments={attachments}
-              setAttachments={setAttachments}
-              messages={messages}
-              setMessages={setMessages}
-              append={append}
-              selectedVisibilityType={visibilityType}
-            />
+            <div className="flex gap-2">
+              <GoogleDrivePicker onFileSelect={handleFileSelect} />
+              <MultimodalInput
+                chatId={id}
+                input={input}
+                setInput={setInput}
+                handleSubmit={handleSubmit}
+                status={status}
+                stop={stop}
+                attachments={attachments}
+                setAttachments={setAttachments}
+                messages={messages}
+                setMessages={setMessages}
+                append={append}
+                selectedVisibilityType={visibilityType}
+              />
+            </div>
           )}
         </form>
       </div>
@@ -175,6 +188,12 @@ export function Chat({
         isReadonly={isReadonly}
         selectedVisibilityType={visibilityType}
       />
+
+      {selectedFile && (
+        <div className="mt-2 text-sm text-gray-500">
+          Selected file: {selectedFile.name}
+        </div>
+      )}
     </>
   );
 }

components/google-drive-picker.tsx

@@ -0,0 +1,95 @@
+'use client';
+
+import { useState } from 'react';
+import { Button } from './ui/button';
+import { useRouter } from 'next/navigation';
+import { GoogleDriveFile } from '@/lib/google-drive';
+
+interface GoogleDrivePickerProps {
+  onFileSelect: (file: GoogleDriveFile) => void;
+}
+
+export function GoogleDrivePicker({ onFileSelect }: GoogleDrivePickerProps) {
+  const [isLoading, setIsLoading] = useState(false);
+  const [error, setError] = useState<string | null>(null);
+  const router = useRouter();
+
+  const handleAuth = () => {
+    console.log('=== Google Drive Auth Flow Start ===');
+    console.log('1. Button clicked, starting auth process');
+
+    try {
+      setIsLoading(true);
+      setError(null);
+
+      // Direct navigation to the auth endpoint
+      console.log('2. Navigating to /api/google-drive/auth');
+      window.location.href = '/api/google-drive/auth';
+    } catch (err) {
+      console.error('=== Google Drive Auth Flow Error ===', err);
+      setError(err instanceof Error ? err.message : 'Failed to start Google Drive authentication');
+      setIsLoading(false);
+    }
+  };
+
+  return (
+    <div className="flex flex-col items-center gap-4">
+      <Button
+        onClick={handleAuth}
+        disabled={isLoading}
+        className="flex items-center gap-2"
+      >
+        {isLoading ? (
+          'Connecting...'
+        ) : (
+          <>
+            <svg
+              className="w-5 h-5"
+              viewBox="0 0 24 24"
+              fill="currentColor"
+            >
+              <path d="M4.433 22l-3.975-6.875h7.95L4.433 22zm7.95-6.875h7.95L12.383 22H4.433l7.95-6.875zm7.95-6.875h7.95L20.333 22h-7.95l7.95-13.75zM12.383 2L4.433 8.875h7.95L12.383 2zm7.95 0l-7.95 6.875h7.95L20.333 2z" />
+            </svg>
+            Connect Google Drive
+          </>
+        )}
+      </Button>
+      {error && (
+        <p className="text-sm text-red-500">{error}</p>
+      )}
+    </div>
+  );
+}
+
+// Helper function to show file selection dialog
+function showFileSelectionDialog(files: GoogleDriveFile[]): Promise<GoogleDriveFile | null> {
+  return new Promise((resolve) => {
+    const dialog = document.createElement('dialog');
+    dialog.className = 'p-4 rounded-lg shadow-lg';
+    dialog.innerHTML = `
+      <div class="max-h-96 overflow-y-auto">
+        <h3 class="text-lg font-semibold mb-4">Select a file</h3>
+        <div class="space-y-2">
+          ${files.map(file => `
+            <button class="w-full text-left p-2 hover:bg-gray-100 rounded" data-id="${file.id}">
+              ${file.name}
+            </button>
+          `).join('')}
+        </div>
+      </div>
+    `;
+
+    dialog.addEventListener('click', (e) => {
+      const target = e.target as HTMLElement;
+      if (target.tagName === 'BUTTON') {
+        const fileId = target.getAttribute('data-id');
+        const selectedFile = files.find(f => f.id === fileId);
+        dialog.close();
+        resolve(selectedFile || null);
+      }
+    });
+
+    document.body.appendChild(dialog);
+    dialog.showModal();
+  });
+} 

components/message.tsx

@@ -12,7 +12,7 @@ import { MessageActions } from './message-actions';
 import { PreviewAttachment } from './preview-attachment';
 import { Weather } from './weather';
 import equal from 'fast-deep-equal';
-import { cn, sanitizeText } from '@/lib/utils';
+import { cn } from '@/lib/utils';
 import { Button } from './ui/button';
 import { Tooltip, TooltipContent, TooltipTrigger } from './ui/tooltip';
 import { MessageEditor } from './message-editor';
@@ -130,7 +130,7 @@ const PurePreviewMessage = ({
                             message.role === 'user',
                         })}
                       >
-                        <Markdown>{sanitizeText(part.text)}</Markdown>
+                        <Markdown>{part.text}</Markdown>
                       </div>
                     </div>
                   );