diff --git a/app/api/generate/route.ts b/app/api/generate/route.ts
index c034158..12fcd0c 100644
--- a/app/api/generate/route.ts
+++ b/app/api/generate/route.ts
@@ -29,7 +29,7 @@ export async function POST(request: NextRequest) {
         const validatedFiles = benchifyFileSchema.parse(generatedFiles);
 
         // // Repair the generated code using Benchify's API
-        const repairedFiles = await repairCode(validatedFiles);
+        const { repairedFiles, buildOutput } = await repairCode(validatedFiles);
 
         const { sbxId, template, url } = await createSandbox({ files: generatedFiles });
 
@@ -39,7 +39,7 @@ export async function POST(request: NextRequest) {
         return NextResponse.json({
             originalFiles: generatedFiles,
             repairedFiles: repairedFiles,
-            // buildOutput: '',  // We don't get build output from Benchify in our current setup
+            buildOutput: buildOutput,
             previewUrl: url,
         });
     } catch (error) {
diff --git a/app/page.tsx b/app/page.tsx
index 542097f..b46a60e 100644
--- a/app/page.tsx
+++ b/app/page.tsx
@@ -4,8 +4,15 @@
 import { useEffect, useState } from 'react';
 import { PromptForm } from '@/components/ui-builder/prompt-form';
 import { Card, CardContent } from '@/components/ui/card';
+import { benchifyFileSchema } from '@/lib/schemas';
+import { z } from 'zod';
+
 export default function Home() {
-  const [result, setResult] = useState<any>(null);
+  const [result, setResult] = useState<{
+    repairedFiles: z.infer<typeof benchifyFileSchema>;
+    buildOutput: string;
+    previewUrl: string;
+  } | null>(null);
 
   useEffect(() => {
     if (result) {
@@ -22,11 +29,19 @@ export default function Home() {
         <p className="text-lg text-muted-foreground mb-8 text-center">
           Generate UI components with AI and automatically repair issues with Benchify
         </p>
-        <Card className="border-border bg-card">
-          <CardContent className="pt-6">
-            <PromptForm onGenerate={setResult} />
-          </CardContent>
-        </Card>
+        {!result ? (
+          <Card className="border-border bg-card">
+            <CardContent className="pt-6">
+              <PromptForm onGenerate={setResult} />
+            </CardContent>
+          </Card>
+        ) : (
+          <Card className="border-border bg-card">
+            <CardContent className="pt-6">
+              <iframe title="Preview" src={result.previewUrl} className="w-full h-full" />
+            </CardContent>
+          </Card>
+        )}
       </div>
     </main>
   );