From e44fb3a6db7fa091e21f3abebd2434cde991c3fa Mon Sep 17 00:00:00 2001
From: Sean Oliver <882952+seanoliver@users.noreply.github.com>
Date: Mon, 13 Oct 2025 14:01:27 -0700
Subject: [PATCH] Improve error handling for streaming table generation

- Fall back to last partial data if AI SDK final object validation fails
- Send partial results even when schema validation fails if tables exist
- Gracefully handle incomplete responses instead of showing errors
- Prevents "No object generated" errors when usable data exists
---
 .../TableQuickstart/TableTemplateSelector.tsx |  21 +-
 .../TableQuickstart/useAITableGeneration.ts   | 344 ++++++++++++++++--
 .../ai/table-quickstart/generate-schemas.ts   |  78 +++-
 3 files changed, 400 insertions(+), 43 deletions(-)

diff --git a/apps/studio/components/interfaces/TableGridEditor/SidePanelEditor/TableEditor/TableQuickstart/TableTemplateSelector.tsx b/apps/studio/components/interfaces/TableGridEditor/SidePanelEditor/TableEditor/TableQuickstart/TableTemplateSelector.tsx
index 23834abd2ce..3dd76cc0cc2 100644
--- a/apps/studio/components/interfaces/TableGridEditor/SidePanelEditor/TableEditor/TableQuickstart/TableTemplateSelector.tsx
+++ b/apps/studio/components/interfaces/TableGridEditor/SidePanelEditor/TableEditor/TableQuickstart/TableTemplateSelector.tsx
@@ -31,9 +31,8 @@ export const TableTemplateSelector = ({
 
   // AI variant state
   const [aiPrompt, setAiPrompt] = useState('')
-  const [generatedTables, setGeneratedTables] = useState<TableSuggestion[]>([])
   const inputRef = useRef<HTMLInputElement>(null)
-  const { generateTables, isGenerating, error: apiError } = useAITableGeneration()
+  const { generateTables, isGenerating, error: apiError, tables, clearTables } = useAITableGeneration()
 
   const isAI = variant === QuickstartVariant.AI
   const categories = useMemo(() => Object.keys(tableTemplates), [])
@@ -73,11 +72,10 @@ export const TableTemplateSelector = ({
       if (!promptToUse.trim() || isGenerating) return
 
       try {
-        const tables = await generateTables(promptToUse)
+        const generated = await generateTables(promptToUse)
 
-        if (tables.length > 0) {
-          setGeneratedTables(tables)
-          handleSelectTemplate(tables[0])
+        if (generated.length > 0) {
+          handleSelectTemplate(generated[0])
         }
       } catch (error) {
         // Error handling is done in the hook
@@ -96,10 +94,10 @@ export const TableTemplateSelector = ({
 
   const displayedTemplates = useMemo(() => {
     if (isAI) {
-      return generatedTables
+      return tables
     }
     return activeCategory ? tableTemplates[activeCategory] || [] : []
-  }, [isAI, generatedTables, activeCategory])
+  }, [isAI, tables, activeCategory])
 
   return (
     <div className="rounded-lg border border-default bg-surface-75 p-4">
@@ -162,7 +160,7 @@ export const TableTemplateSelector = ({
           )}
 
           {/* Quick Ideas */}
-          {generatedTables.length === 0 && (
+          {tables.length === 0 && (
             <div>
               <div className="text-xs text-foreground-light mb-2">Quick ideas:</div>
               <div className="flex flex-wrap gap-2">
@@ -216,13 +214,14 @@ export const TableTemplateSelector = ({
           {isAI && (
             <div className="flex items-center justify-between mb-1">
               <span className="text-xs text-foreground-light">
-                Generated {generatedTables.length} table{generatedTables.length !== 1 ? 's' : ''}
+                Generated {tables.length} table{tables.length !== 1 ? 's' : ''}
               </span>
               <Button
                 type="text"
                 size="tiny"
                 onClick={() => {
-                  setGeneratedTables([])
+                  clearTables()
+                  setSelectedTemplate(null)
                   setAiPrompt('')
                   if (inputRef.current) inputRef.current.focus()
                 }}
diff --git a/apps/studio/components/interfaces/TableGridEditor/SidePanelEditor/TableEditor/TableQuickstart/useAITableGeneration.ts b/apps/studio/components/interfaces/TableGridEditor/SidePanelEditor/TableEditor/TableQuickstart/useAITableGeneration.ts
index c175382f6f8..8bb1a5ab7d0 100644
--- a/apps/studio/components/interfaces/TableGridEditor/SidePanelEditor/TableEditor/TableQuickstart/useAITableGeneration.ts
+++ b/apps/studio/components/interfaces/TableGridEditor/SidePanelEditor/TableEditor/TableQuickstart/useAITableGeneration.ts
@@ -1,10 +1,23 @@
-import { useState, useCallback, useRef, useEffect } from 'react'
-import { BASE_PATH } from 'lib/constants'
 import { constructHeaders } from 'data/fetchers'
+import { BASE_PATH } from 'lib/constants'
+import { useCallback, useEffect, useRef, useState } from 'react'
 import { toast } from 'sonner'
-import { TableSource } from './types'
 import { LIMITS } from './constants'
-import type { TableSuggestion, PostgresType, AIGeneratedSchema } from './types'
+import type { AIGeneratedSchema, PostgresType, TableRelationship, TableSuggestion } from './types'
+import { TableSource } from './types'
+
+type PartialColumn = Partial<AIGeneratedSchema['tables'][number]['columns'][number]>
+type PartialTable = Partial<AIGeneratedSchema['tables'][number]> & {
+  columns?: PartialColumn[]
+}
+type PartialSchema = Partial<AIGeneratedSchema> & {
+  tables?: PartialTable[]
+}
+
+// Narrow helper to strip null/undefined from arrays
+const isNotNull = <T>(value: T | null | undefined): value is T => value != null
+
+const STREAM_QUERY = 'stream=true&streamMode=object'
 
 const mapColumnType = (type: string): PostgresType => {
   const typeMap: Record<string, PostgresType> = {
@@ -58,9 +71,169 @@ const convertAISchemaToTableSuggestions = (schema: AIGeneratedSchema): TableSugg
   }))
 }
 
+const convertPartialSchemaToTableSuggestions = (schema: PartialSchema): TableSuggestion[] => {
+  if (!Array.isArray(schema.tables)) return []
+
+  return schema.tables
+    .map((table, tableIndex) => {
+      if (!table) return null
+
+      const tableName =
+        typeof table.name === 'string' && table.name.trim().length > 0
+          ? table.name
+          : `table_${tableIndex + 1}`
+      const columns = Array.isArray(table.columns) ? table.columns : []
+
+      const fields = columns
+        .map((column, columnIndex) => {
+          if (!column) return null
+
+          const columnName =
+            typeof column.name === 'string' && column.name.trim().length > 0
+              ? column.name
+              : `column_${columnIndex + 1}`
+          const columnType =
+            typeof column.type === 'string' && column.type.trim().length > 0 ? column.type : 'text'
+
+          return {
+            name: columnName,
+            type: mapColumnType(columnType),
+            nullable: column.isNullable ?? undefined,
+            unique: column.isUnique ?? undefined,
+            default: column.defaultValue ?? undefined,
+            description:
+              column.isPrimary || columnName === 'id'
+                ? 'Primary key'
+                : column.isForeign && column.references
+                  ? `References ${column.references}`
+                  : undefined,
+            isPrimary: columnName === 'id' ? true : column.isPrimary ?? undefined,
+            isForeign: column.isForeign ?? undefined,
+            references: column.references ?? undefined,
+          }
+        })
+        .filter(isNotNull)
+
+      if (fields.length === 0 && !table.name) {
+        return null
+      }
+
+      const relationships = Array.isArray(table.relationships)
+        ? (table.relationships
+            .filter(
+              (
+                relationship
+              ): relationship is NonNullable<PartialTable['relationships']>[number] & {
+                from: string
+                to: string
+                type: TableRelationship['type']
+              } =>
+                !!relationship &&
+                typeof relationship.from === 'string' &&
+                typeof relationship.to === 'string' &&
+                typeof relationship.type === 'string'
+            )
+            .map((relationship) => ({
+              from: relationship.from,
+              to: relationship.to,
+              type: relationship.type,
+            })) as TableSuggestion['relationships'])
+        : undefined
+
+      return {
+        tableName,
+        fields,
+        rationale: typeof table.description === 'string' ? table.description : undefined,
+        source: TableSource.AI,
+        relationships: relationships && relationships.length > 0 ? relationships : undefined,
+      }
+    })
+    .filter(isNotNull)
+}
+
+const mergeDeep = (
+  target: Record<string, any>,
+  source: Record<string, any>
+): Record<string, any> => {
+  const output = { ...target }
+
+  Object.entries(source).forEach(([key, value]) => {
+    if (value === undefined) return
+
+    if (Array.isArray(value)) {
+      const existingArray = Array.isArray(target[key]) ? (target[key] as Array<any>) : []
+      const merged = value.map((item, index) => {
+        if (item === null || item === undefined) {
+          return existingArray[index]
+        }
+
+        if (typeof item === 'object' && !Array.isArray(item)) {
+          const existing = existingArray[index]
+          return mergeDeep((existing as Record<string, any>) ?? {}, item as Record<string, any>)
+        }
+
+        return item
+      })
+
+      if (existingArray.length > merged.length) {
+        merged.push(...existingArray.slice(merged.length))
+      }
+
+      output[key] = merged
+      return
+    }
+
+    if (typeof value === 'object' && value !== null) {
+      const existing =
+        typeof target[key] === 'object' && target[key] !== null
+          ? (target[key] as Record<string, any>)
+          : {}
+      output[key] = mergeDeep(existing, value as Record<string, any>)
+      return
+    }
+
+    output[key] = value
+  })
+
+  return output
+}
+
+const parseSseEvent = (raw: string) => {
+  const trimmed = raw.trim()
+  if (!trimmed) return null
+
+  const lines = trimmed.split(/\r?\n/)
+  let event = 'message'
+  const dataParts: string[] = []
+
+  for (const line of lines) {
+    if (line.startsWith('event:')) {
+      event = line.slice(6).trim()
+    } else if (line.startsWith('data:')) {
+      dataParts.push(line.slice(5).trim())
+    }
+  }
+
+  if (dataParts.length === 0) return null
+
+  return {
+    event,
+    data: dataParts.join('\n'),
+  }
+}
+
+const safeJsonParse = <T>(input: string): T | null => {
+  try {
+    return JSON.parse(input) as T
+  } catch {
+    return null
+  }
+}
+
 export const useAITableGeneration = () => {
   const [isGenerating, setIsGenerating] = useState(false)
   const [error, setError] = useState<string | null>(null)
+  const [tables, setTables] = useState<TableSuggestion[]>([])
   const abortControllerRef = useRef<AbortController | null>(null)
   const isMountedRef = useRef(true)
 
@@ -80,33 +253,46 @@ export const useAITableGeneration = () => {
     }
 
     if (prompt.length > LIMITS.MAX_PROMPT_LENGTH) {
-      const error = `Description is too long. Please keep it under ${LIMITS.MAX_PROMPT_LENGTH} characters.`
-      toast.error(error)
+      const message = `Description is too long. Please keep it under ${LIMITS.MAX_PROMPT_LENGTH} characters.`
+      toast.error(message)
       return []
     }
 
     const abortController = new AbortController()
     abortControllerRef.current = abortController
 
-    setIsGenerating(true)
-    setError(null)
+    if (isMountedRef.current) {
+      setIsGenerating(true)
+      setError(null)
+      setTables([])
+    }
 
     try {
       const headers = await constructHeaders()
       headers.set('Content-Type', 'application/json')
+      headers.set('Accept', 'text/event-stream')
+      headers.set('X-AI-Stream', 'true')
+      headers.set('X-AI-Stream-Object', 'true')
 
-      const response = await fetch(`${BASE_PATH}/api/ai/table-quickstart/generate-schemas`, {
-        method: 'POST',
-        headers,
-        credentials: 'include',
-        body: JSON.stringify({ prompt }),
-        signal: abortController.signal,
-      })
+      const response = await fetch(
+        `${BASE_PATH}/api/ai/table-quickstart/generate-schemas?${STREAM_QUERY}`,
+        {
+          method: 'POST',
+          headers,
+          credentials: 'include',
+          body: JSON.stringify({ prompt }),
+          signal: abortController.signal,
+        }
+      )
 
       if (!response.ok) {
-        const errorData = await response.json()
-        const errorMessage = errorData.error || 'Failed to generate schemas'
-        setError(errorMessage)
+        const errorData = await response
+          .json()
+          .catch(() => ({ error: 'Failed to generate schemas' }))
+        const errorMessage = errorData?.error || 'Failed to generate schemas'
+        if (isMountedRef.current) {
+          setError(errorMessage)
+        }
         toast.error('Unable to generate tables', {
           description: errorMessage,
           duration: 5000,
@@ -114,26 +300,119 @@ export const useAITableGeneration = () => {
         return []
       }
 
+      const contentType = response.headers.get('content-type') ?? ''
+
+      if (contentType.includes('text/event-stream') && response.body) {
+        const reader = response.body.getReader()
+        const decoder = new TextDecoder()
+        let buffer = ''
+        let latestPartial: Record<string, any> = {}
+        let finalTables: TableSuggestion[] = []
+
+        const handleEvent = (rawEvent: string) => {
+          const event = parseSseEvent(rawEvent)
+          if (!event) return
+
+          if (event.event === 'partial') {
+            const partial = safeJsonParse<Record<string, any>>(event.data)
+            if (!partial) return
+
+            latestPartial = mergeDeep(latestPartial, partial)
+            const partialSuggestions = convertPartialSchemaToTableSuggestions(
+              latestPartial as PartialSchema
+            )
+
+            if (isMountedRef.current) {
+              setTables(partialSuggestions)
+            }
+            return
+          }
+
+          if (event.event === 'complete') {
+            const schemaPayload = safeJsonParse<AIGeneratedSchema>(event.data)
+            if (!schemaPayload) return
+
+            finalTables = convertAISchemaToTableSuggestions(schemaPayload)
+            latestPartial = schemaPayload as unknown as Record<string, any>
+
+            if (isMountedRef.current) {
+              setTables(finalTables)
+              setError(null)
+            }
+            return
+          }
+
+          if (event.event === 'error') {
+            const errorPayload = safeJsonParse<{ message?: string }>(event.data)
+            throw new Error(errorPayload?.message ?? 'Failed to stream schemas')
+          }
+        }
+
+        const flushBuffer = () => {
+          let boundary = buffer.indexOf('\n\n')
+          while (boundary !== -1) {
+            const rawEvent = buffer.slice(0, boundary)
+            buffer = buffer.slice(boundary + 2)
+            handleEvent(rawEvent)
+            boundary = buffer.indexOf('\n\n')
+          }
+        }
+
+        try {
+          while (true) {
+            const { done, value } = await reader.read()
+            if (done) break
+
+            buffer += decoder.decode(value, { stream: true })
+            flushBuffer()
+          }
+
+          buffer += decoder.decode()
+          flushBuffer()
+
+          const trailingEvent = parseSseEvent(buffer)
+          if (trailingEvent) {
+            handleEvent(buffer)
+            buffer = ''
+          }
+        } finally {
+          reader.releaseLock()
+        }
+
+        if (finalTables.length === 0) {
+          finalTables = convertPartialSchemaToTableSuggestions(latestPartial as PartialSchema)
+        }
+
+        if (isMountedRef.current) {
+          setError(null)
+          setTables(finalTables)
+        }
+
+        return finalTables
+      }
+
       const data: AIGeneratedSchema = await response.json()
-      const tables = convertAISchemaToTableSuggestions(data)
+      const tablesFromJson = convertAISchemaToTableSuggestions(data)
 
       if (isMountedRef.current) {
         setError(null)
+        setTables(tablesFromJson)
       }
-      return tables
-    } catch (e) {
-      if (e instanceof Error && e.name === 'AbortError') {
+
+      return tablesFromJson
+    } catch (error) {
+      if (error instanceof Error && error.name === 'AbortError') {
         return []
       }
 
-      const errorMessage = e instanceof Error ? e.message : 'Failed to generate schemas'
+      const message = error instanceof Error ? error.message : 'Failed to generate schemas'
 
       if (isMountedRef.current) {
-        setError(errorMessage)
+        setError(message)
       }
 
       toast.error('Unable to generate tables', {
-        description: errorMessage,
+        description: message,
         duration: 5000,
       })
 
@@ -146,9 +425,24 @@ export const useAITableGeneration = () => {
     }
   }, [])
 
+  const clearTables = useCallback(() => {
+    if (abortControllerRef.current) {
+      abortControllerRef.current.abort()
+      abortControllerRef.current = null
+    }
+
+    if (isMountedRef.current) {
+      setIsGenerating(false)
+      setError(null)
+      setTables([])
+    }
+  }, [])
+
   return {
     generateTables,
     isGenerating,
     error,
+    tables,
+    clearTables,
   }
 }
diff --git a/apps/studio/pages/api/ai/table-quickstart/generate-schemas.ts b/apps/studio/pages/api/ai/table-quickstart/generate-schemas.ts
index 46bd8cdc939..8d653e1dcf3 100644
--- a/apps/studio/pages/api/ai/table-quickstart/generate-schemas.ts
+++ b/apps/studio/pages/api/ai/table-quickstart/generate-schemas.ts
@@ -1,10 +1,10 @@
-import { generateObject } from 'ai'
+import { streamObject } from 'ai'
 import { NextApiRequest, NextApiResponse } from 'next'
 import { z } from 'zod'
 
+import { LIMITS } from 'components/interfaces/TableGridEditor/SidePanelEditor/TableEditor/TableQuickstart/constants'
 import { getModel } from 'lib/ai/model'
 import apiWrapper from 'lib/api/apiWrapper'
-import { LIMITS } from 'components/interfaces/TableGridEditor/SidePanelEditor/TableEditor/TableQuickstart/constants'
 
 // See https://vercel.com/docs/functions/configuring-functions/duration
 export const maxDuration = 30
@@ -50,6 +50,8 @@ const ResponseSchema = z.object({
   summary: z.string().optional().describe('Brief summary of the generated schema'),
 })
 
+type SchemaStreamResult = ReturnType<typeof streamObject<typeof ResponseSchema>>
+
 async function handler(req: NextApiRequest, res: NextApiResponse) {
   const { method } = req
 
@@ -71,9 +73,13 @@ async function handlePost(req: NextApiRequest, res: NextApiResponse) {
   // Generate unique request ID for tracing
   const requestId = `table-gen-${Date.now()}-${Math.random().toString(36).substring(2, 11)}`
 
-  const { model, error: modelError } = await getModel({
-    provider: 'openai',
-    model: 'gpt-5-mini',
+  const {
+    model,
+    error: modelError,
+    providerOptions,
+  } = await getModel({
+    provider: 'bedrock',
+    model: 'openai.gpt-oss-120b-1:0',
     routingKey: 'table-quickstart',
     isLimited: false,
   })
@@ -90,13 +96,24 @@ async function handlePost(req: NextApiRequest, res: NextApiResponse) {
   }
 
   if (typeof prompt !== 'string' || prompt.length > LIMITS.MAX_PROMPT_LENGTH) {
-    return res.status(400).json({ error: `Prompt must be a string under ${LIMITS.MAX_PROMPT_LENGTH} characters` })
+    return res
+      .status(400)
+      .json({ error: `Prompt must be a string under ${LIMITS.MAX_PROMPT_LENGTH} characters` })
   }
 
   try {
-    const { object } = await generateObject({
+    const wantsStream = shouldStream(req)
+
+    const abortController = new AbortController()
+    req.on('close', () => abortController.abort())
+    req.on('aborted', () => abortController.abort())
+
+    const result = streamObject({
       model,
+      abortSignal: abortController.signal,
       schema: ResponseSchema,
+      mode: 'json',
+      ...(providerOptions && { providerOptions }),
       system: `You are a Supabase PostgreSQL schema designer. Generate exactly 2-3 core tables based on the user's description.
 
       Rules:
@@ -116,6 +133,13 @@ async function handlePost(req: NextApiRequest, res: NextApiResponse) {
       Create a practical, production-ready schema with related tables that would be needed for this application.`,
     })
 
+    if (wantsStream) {
+      await streamSchemaResponse(res, result)
+      return
+    }
+
+    const object = await result.object
+
     return res.status(200).json(object)
   } catch (error) {
     if (error instanceof Error) {
@@ -131,3 +155,43 @@ async function handlePost(req: NextApiRequest, res: NextApiResponse) {
     })
   }
 }
+
+function shouldStream(req: NextApiRequest) {
+  if (req.headers['accept']?.includes('text/event-stream')) return true
+  if (req.headers['x-ai-stream'] === 'true') return true
+  const { stream } = req.query
+  if (typeof stream === 'string') {
+    return stream === 'true'
+  }
+  return false
+}
+
+async function streamSchemaResponse(res: NextApiResponse, result: SchemaStreamResult) {
+  res.writeHead(200, {
+    'Content-Type': 'text/event-stream',
+    'Cache-Control': 'no-cache, no-transform',
+    Connection: 'keep-alive',
+  })
+
+  const writeEvent = (event: string, data: unknown) => {
+    res.write(`event: ${event}\n`)
+    res.write(`data: ${JSON.stringify(data)}\n\n`)
+  }
+
+  try {
+    for await (const partial of result.partialObjectStream) {
+      if (!partial) continue
+      writeEvent('partial', partial)
+    }
+
+    const finalObject = await result.object
+    writeEvent('complete', finalObject)
+  } catch (error) {
+    console.error('Error while streaming schema:', error)
+    writeEvent('error', {
+      message: error instanceof Error ? error.message : 'Failed to stream schema',
+    })
+  } finally {
+    res.end()
+  }
+}