Daily branch 2025 01 10 (#807)

* fix(mfa-login): server components render

* fix(mfa-login): error handeling

* refactor(tools-handler): use gpt-4o-mini for free users

* refactor(mistral): remove @ai-sdk/mistral and code related

* fix: improve auth callback handling in middleware

* fix: improve auth callback handling in middleware

* revert: improve auth callback handling in middleware

* feat: add 'use server' directive to checkAuth function

* feat: add mfa check on account delete

* feat: add mfa check on account delete

app/api/chat/mistral/route.ts

@@ -11,8 +11,7 @@ import { handleErrorResponse } from "@/lib/models/llm/api-error"
 import llmConfig from "@/lib/models/llm/llm-config"
 import { generateStandaloneQuestion } from "@/lib/models/question-generator"
 import { checkRatelimitOnApi } from "@/lib/server/ratelimiter"
-import { createMistral } from "@ai-sdk/mistral"
-import { createOpenAI } from "@ai-sdk/openai"
+import { createOpenAI as createOpenRouterAI } from "@ai-sdk/openai"
 import { createDataStreamResponse, streamText } from "ai"
 import { getModerationResult } from "@/lib/server/moderation"
 import { createToolSchemas } from "@/lib/tools/llm/toolSchemas"
@@ -63,7 +62,6 @@ export async function POST(request: Request) {
     let {
       providerBaseUrl,
       providerHeaders,
-      providerApiKey,
       selectedModel,
       rateLimitCheckResult,
       similarityTopK,
@@ -179,24 +177,10 @@ export async function POST(request: Request) {
     handleMessages(shouldUncensorResponse)
 
     try {
-      let provider
-
-      if (selectedModel.startsWith("mistralai")) {
-        provider = createMistral({
-          apiKey: providerApiKey,
-          baseURL: providerBaseUrl,
-          headers: providerHeaders
-        })
-      } else if (selectedModel.startsWith("mistral")) {
-        provider = createMistral({
-          apiKey: llmConfig.mistral.apiKey
-        })
-      } else {
-        provider = createOpenAI({
-          baseURL: providerBaseUrl,
-          headers: providerHeaders
-        })
-      }
+      const provider = createOpenRouterAI({
+        baseURL: providerBaseUrl,
+        headers: providerHeaders
+      })
 
       // Handle web search plugin
       switch (selectedPlugin) {
@@ -306,7 +290,6 @@ async function getProviderConfig(
   return {
     providerUrl,
     providerBaseUrl,
-    providerApiKey,
     providerHeaders,
     selectedModel,
     rateLimitCheckResult,

app/login/page.tsx

@@ -81,6 +81,8 @@ export default async function Login({
   }
 
   const checkAuth = async () => {
+    "use server"
+
     const supabase = await createClient()
     const {
       data: { user }

app/login/verify/mfa-verification.tsx

@@ -11,19 +11,20 @@ import { PentestGPTContext } from "@/context/context"
 import { getHomeWorkspaceByUserId } from "@/db/workspaces"
 
 interface MFAVerificationProps {
-  onVerify: (code: string) => Promise<{ success: boolean } | void>
+  onVerify: (code: string) => Promise<{ success: boolean; error?: string }>
 }
 
 export function MFAVerification({ onVerify }: MFAVerificationProps) {
   const router = useRouter()
   const [verifyCode, setVerifyCode] = useState("")
   const [error, setError] = useState("")
   const [isVerifying, setIsVerifying] = useState(false)
-  const { user } = useContext(PentestGPTContext)
-  const { fetchStartingData } = useContext(PentestGPTContext)
+  const { user, fetchStartingData } = useContext(PentestGPTContext)
 
   const handleSubmit = async (e: React.FormEvent) => {
     e.preventDefault()
+    if (!user || verifyCode.length !== 6 || isVerifying) return
+
     setError("")
     setIsVerifying(true)
 
@@ -38,60 +39,27 @@ export function MFAVerification({ onVerify }: MFAVerificationProps) {
         await fetchStartingData()
         const homeWorkspaceId = await getHomeWorkspaceByUserId(user.id)
         router.push(`/${homeWorkspaceId}/chat`)
-      }
-    } catch (err) {
-      setIsVerifying(false)
-
-      // Handle Supabase Auth specific errors
-      if (err instanceof Error) {
-        const error = err as any // for accessing .code property
-
-        switch (error.code) {
-          case "mfa_verification_failed":
-            setError("Invalid verification code. Please try again.")
-            break
-          case "mfa_challenge_expired":
-            setError("Verification code has expired. Please request a new one.")
-            break
-          case "over_request_rate_limit":
-            setError(
-              "Too many attempts. Please wait a few minutes before trying again."
-            )
-            break
-          case "mfa_totp_verify_not_enabled":
-            setError(
-              "MFA verification is currently disabled. Please contact support."
-            )
-            break
-          case "mfa_verification_rejected":
-            setError(
-              "Verification was rejected. Please try again or contact support."
-            )
-            break
-          default:
-            // Check for specific error messages as fallback
-            if (error.message?.includes("Invalid one-time password")) {
-              setError("Invalid verification code. Please try again.")
-            } else if (error.message?.includes("rate limit")) {
-              setError("Too many attempts. Please wait a moment and try again.")
-            } else {
-              console.error("MFA Verification Error:", error)
-              setError(
-                "Unable to verify code. Please try again or contact support."
-              )
-            }
-        }
       } else {
-        console.error("Unknown MFA Error:", err)
-        setError("An unexpected error occurred. Please try again.")
+        setError(result.error || "Verification failed")
+        setVerifyCode("")
       }
+    } catch (error) {
+      setError("Please try again")
+      setVerifyCode("")
+    } finally {
+      setIsVerifying(false)
     }
   }
 
   const handleSignOut = async () => {
-    await supabase.auth.signOut({ scope: "local" })
-    router.push("/login")
-    router.refresh()
+    try {
+      await supabase.auth.signOut({ scope: "local" })
+      router.push("/login")
+      router.refresh()
+    } catch (error) {
+      console.error("Sign out error:", error)
+      router.push("/login")
+    }
   }
 
   return (

app/login/verify/page.tsx

@@ -2,6 +2,11 @@ import { createClient } from "@/lib/supabase/server"
 import { redirect } from "next/navigation"
 import { MFAVerification } from "./mfa-verification"
 
+interface VerifyMFAResponse {
+  success: boolean
+  error?: string
+}
+
 export default async function VerifyMFA() {
   const supabase = await createClient()
 
@@ -25,37 +30,53 @@ export default async function VerifyMFA() {
     return redirect("/login")
   }
 
-  const verifyMFA = async (code: string): Promise<{ success: boolean }> => {
+  const verifyMFA = async (code: string): Promise<VerifyMFAResponse> => {
     "use server"
 
     const supabase = await createClient()
 
     try {
-      const { data: factors, error: factorsError } =
-        await supabase.auth.mfa.listFactors()
-      if (factorsError) throw factorsError
+      const { data: factors } = await supabase.auth.mfa.listFactors()
+      const totpFactor = factors?.totp[0]
 
-      const totpFactor = factors.totp[0]
       if (!totpFactor) {
-        throw new Error("No TOTP factors found!")
+        return {
+          success: false,
+          error: "MFA verification is not set up properly"
+        }
       }
 
-      const { data: challenge, error: challengeError } =
-        await supabase.auth.mfa.challenge({
-          factorId: totpFactor.id
-        })
-      if (challengeError) throw challengeError
+      const { data: challenge } = await supabase.auth.mfa.challenge({
+        factorId: totpFactor.id
+      })
+
+      if (!challenge) {
+        return {
+          success: false,
+          error: "Failed to initiate verification"
+        }
+      }
 
       const { error: verifyError } = await supabase.auth.mfa.verify({
         factorId: totpFactor.id,
         challengeId: challenge.id,
         code
       })
-      if (verifyError) throw verifyError
+
+      if (verifyError) {
+        return {
+          success: false,
+          error: "Invalid verification code"
+        }
+      }
 
       return { success: true }
     } catch (error) {
-      throw error
+      console.error("MFA verification error:", error)
+      return {
+        success: false,
+        error: "Verification failed. Please try again"
+      }
     }
   }