Fix chat persistence and sync Supabase credits migration

app/actions/generate.ts

@@ -23,6 +23,23 @@ type GenerateProps = {
   researchMode?: boolean
 }
 
+function isMissingColumnError(error: unknown, columnName: string) {
+  if (!error || typeof error !== 'object') {
+    return false
+  }
+
+  const details = [
+    'message' in error ? error.message : '',
+    'details' in error ? error.details : '',
+    'hint' in error ? error.hint : '',
+  ]
+    .filter((value): value is string => typeof value === 'string')
+    .join(' ')
+    .toLowerCase()
+
+  return details.includes(columnName.toLowerCase()) && details.includes('column')
+}
+
 export async function generateAnswer({ prompt, tool, authorId, authorEmail, attachments = [], sessionId, chatId, enableWebSearch = false, researchMode = false }: GenerateProps) {
   // Get user profile and check limits
   let userProfile = await getUserProfileServer(authorId)
@@ -172,19 +189,31 @@ export async function generateAnswer({ prompt, tool, authorId, authorEmail, atta
 
   if (chatId) {
     // Update existing row
-    const { error } = await supabaseServer
+    const baseUpdatePayload = {
+      prompt,
+      response: answer,
+      attachments,
+    }
+
+    let { error } = await supabaseServer
       .from('chat_sessions')
       .update({
-        prompt,
-        response: answer,
-        attachments,
+        ...baseUpdatePayload,
         token_usage: tokenCost,
-        // We don't update created_at or session_id usually, but ensure they match just in case?
-        // Usually just updating content is enough.
       })
       .eq('id', chatId)
       .eq('user_id', authorId)
 
+    if (error && isMissingColumnError(error, 'token_usage')) {
+      const retryResult = await supabaseServer
+        .from('chat_sessions')
+        .update(baseUpdatePayload)
+        .eq('id', chatId)
+        .eq('user_id', authorId)
+
+      error = retryResult.error
+    }
+
     if (error) {
       console.error('[chat_update_failed]', { userId: authorId, chatId, error })
       persistenceWarning = 'We generated your response, but could not save this chat message.'
@@ -193,20 +222,33 @@ export async function generateAnswer({ prompt, tool, authorId, authorEmail, atta
     }
   } else {
     // Insert new row
-    const { data, error } = await supabaseServer.from('chat_sessions').insert({
+    const baseInsertPayload = {
       user_id: authorId,
       tool,
       prompt,
       response: answer,
       attachments,
-      token_usage: tokenCost,
       created_at: new Date().toISOString(),
       session_id: currentSessionId,
       title: title
+    }
+
+    let { data, error } = await supabaseServer.from('chat_sessions').insert({
+      ...baseInsertPayload,
+      token_usage: tokenCost,
     })
       .select('id')
       .single()
 
+    if (error && isMissingColumnError(error, 'token_usage')) {
+      const retryResult = await supabaseServer.from('chat_sessions').insert(baseInsertPayload)
+        .select('id')
+        .single()
+
+      data = retryResult.data
+      error = retryResult.error
+    }
+
     if (error) {
       console.error('[chat_insert_failed]', { userId: authorId, sessionId: currentSessionId, error })
       persistenceWarning = 'We generated your response, but could not save this chat message.'

components/AuthProvider.tsx

@@ -20,7 +20,7 @@ const AuthContext = createContext<AuthContextType | undefined>(undefined)
 function AuthContextProvider({ children }: { children: React.ReactNode }) {
   const { data: session, status } = useSession()
   const loading = status === 'loading'
-  const userReady = status === 'authenticated' && !!session?.user
+  const userReady = status === 'authenticated' && !!session?.user?.id
 
   const user = session?.user ? {
     id: session.user.id as string,

lib/auth.ts

@@ -95,11 +95,42 @@ export const { handlers, signIn, signOut, auth } = NextAuth((req) => {
           token.name = user.name
           token.picture = user.image
         }
+
+        if (!token.userId && token.email) {
+          try {
+            const { data } = await supabaseServer
+              .from('users')
+              .select('id')
+              .eq('email', token.email)
+              .maybeSingle()
+
+            if (data?.id) {
+              token.userId = data.id
+            }
+          } catch {
+          }
+        }
+
         return token
       },
 
       async session({ session, token }) {
         if (session.user && token) {
+          if (!token.userId && token.email) {
+            try {
+              const { data } = await supabaseServer
+                .from('users')
+                .select('id')
+                .eq('email', token.email)
+                .maybeSingle()
+
+              if (data?.id) {
+                token.userId = data.id
+              }
+            } catch {
+            }
+          }
+
           session.user.id = token.userId as string
           session.user.email = token.email as string
           session.user.name = token.name as string

lib/free-plan-credits.ts

@@ -23,6 +23,23 @@ type UserCreditRecord = {
   resetDate: Date | null
 }
 
+function isMissingColumnError(error: unknown, columnName: string) {
+  if (!error || typeof error !== 'object') {
+    return false
+  }
+
+  const details = [
+    'message' in error ? error.message : '',
+    'details' in error ? error.details : '',
+    'hint' in error ? error.hint : '',
+  ]
+    .filter((value): value is string => typeof value === 'string')
+    .join(' ')
+    .toLowerCase()
+
+  return details.includes(columnName.toLowerCase()) && details.includes('column')
+}
+
 export function getFreePlanCreditCap() {
   return PLAN_MONTHLY_CREDIT_CAPS.free
 }
@@ -52,6 +69,24 @@ async function getUserCreditRecord(userId: string): Promise<UserCreditRecord | n
     .eq('id', userId)
     .maybeSingle()
 
+  if (error && (isMissingColumnError(error, 'free_plan_credits_used') || isMissingColumnError(error, 'free_plan_credits_reset_date'))) {
+    const { data: fallbackData, error: fallbackError } = await supabaseServer
+      .from('users')
+      .select('subscription_plan')
+      .eq('id', userId)
+      .maybeSingle()
+
+    if (fallbackError || !fallbackData) {
+      return null
+    }
+
+    return {
+      plan: normalizePlan(fallbackData.subscription_plan),
+      used: 0,
+      resetDate: null,
+    }
+  }
+
   if (error || !data) {
     return null
   }
@@ -84,6 +119,20 @@ async function getSessionCreditUsage(userId: string, windowStart: Date): Promise
     .eq('user_id', userId)
     .gte('created_at', windowStart.toISOString())
 
+  if (error && isMissingColumnError(error, 'token_usage')) {
+    const { data: fallbackData, error: fallbackError } = await supabaseServer
+      .from('chat_sessions')
+      .select('id')
+      .eq('user_id', userId)
+      .gte('created_at', windowStart.toISOString())
+
+    if (fallbackError || !fallbackData) {
+      return 0
+    }
+
+    return fallbackData.length
+  }
+
   if (error || !data) {
     return 0
   }