stream PUT request to avoid timeout from cloudfront

app/api/ask/route.ts

@@ -6,35 +6,19 @@ import { InferenceClient } from "@huggingface/inference";
 
 import { MODELS } from "@/lib/providers";
 import {
-  DIVIDER,
   FOLLOW_UP_SYSTEM_PROMPT,
   INITIAL_SYSTEM_PROMPT,
   MAX_REQUESTS_PER_IP,
-  NEW_FILE_END,
-  NEW_FILE_START,
-  REPLACE_END,
-  SEARCH_START,
-  UPDATE_FILE_START,
-  UPDATE_FILE_END,
   PROMPT_FOR_PROJECT_NAME,
 } from "@/lib/prompts";
 import { calculateMaxTokens, estimateInputTokens, getProviderSpecificConfig } from "@/lib/max-tokens";
 import MY_TOKEN_KEY from "@/lib/get-cookie-name";
 import { Page } from "@/types";
-import { createRepo, RepoDesignation, uploadFiles } from "@huggingface/hub";
 import { isAuthenticated } from "@/lib/auth";
 import { getBestProvider } from "@/lib/best-provider";
-// import { rewritePrompt } from "@/lib/rewrite-prompt";
-import { COLORS } from "@/lib/utils";
-import { templates } from "@/lib/templates";
-import { injectDeepSiteBadge, isIndexPage } from "@/lib/inject-badge";
 
 const ipAddresses = new Map();
 
-const STREAMING_TIMEOUT = 180000;
-const REQUEST_TIMEOUT = 240000;
-export const maxDuration = 300;
-
 export async function POST(request: NextRequest) {
   const authHeaders = await headers();
   const userToken = request.cookies.get(MY_TOKEN_KEY())?.value;
@@ -117,9 +101,6 @@ export async function POST(request: NextRequest) {
 
     (async () => {
       // let completeResponse = "";
-      let timeoutId: NodeJS.Timeout | null = null;
-      let isTimedOut = false;
-      
       try {
         const client = new InferenceClient(token);
         
@@ -151,51 +132,21 @@ export async function POST(request: NextRequest) {
           billTo ? { billTo } : {}
         );
 
-        // Set up timeout
-        const timeoutPromise = new Promise((_, reject) => {
-          timeoutId = setTimeout(() => {
-            isTimedOut = true;
-            reject(new Error('Request timeout: The AI model took too long to respond. Please try again with a simpler prompt or try a different model.'));
-          }, STREAMING_TIMEOUT);
-        });
-
-        // Race between streaming and timeout
-        await Promise.race([
-          (async () => {
-            while (true) {
-              const { done, value } = await chatCompletion.next()
-              if (done) {
-                break;
-              }
-
-              const chunk = value.choices[0]?.delta?.content;
-              if (chunk) {
-                await writer.write(encoder.encode(chunk));
-              }
-            }
-          })(),
-          timeoutPromise
-        ]);
-
-        // Clear timeout if successful
-        if (timeoutId) clearTimeout(timeoutId);
+        while (true) {
+          const { done, value } = await chatCompletion.next()
+          if (done) {
+            break;
+          }
+
+          const chunk = value.choices[0]?.delta?.content;
+          if (chunk) {
+            await writer.write(encoder.encode(chunk));
+          }
+        }
         
-        // Explicitly close the writer after successful completion
         await writer.close();
       } catch (error: any) {
-        // Clear timeout on error
-        if (timeoutId) clearTimeout(timeoutId);
-        
-        if (isTimedOut || error.message?.includes('timeout') || error.message?.includes('Request timeout')) {
-          await writer.write(
-            encoder.encode(
-              JSON.stringify({
-                ok: false,
-                message: "Request timeout: The AI model took too long to respond. Please try again with a simpler prompt or try a different model.",
-              })
-            )
-          );
-        } else if (error.message?.includes("exceeded your monthly included credits")) {
+        if (error.message?.includes("exceeded your monthly included credits")) {
           await writer.write(
             encoder.encode(
               JSON.stringify({
@@ -259,11 +210,9 @@ export async function PUT(request: NextRequest) {
   const authHeaders = await headers();
 
   const body = await request.json();
-  const { prompt, provider, selectedElementHtml, model, pages, files, repoId: repoIdFromBody, isNew } =
+  const { prompt, provider, selectedElementHtml, model, pages, files, repoId, isNew } =
     body;
 
-  let repoId = repoIdFromBody;
-
   if (!prompt || pages.length === 0) {
     return NextResponse.json(
       { ok: false, error: "Missing required fields" },
@@ -314,453 +263,133 @@ export async function PUT(request: NextRequest) {
     billTo = "huggingface";
   }
 
-  const client = new InferenceClient(token);
+  const selectedProvider = await getBestProvider(selectedModel.value, provider);
 
-  const escapeRegExp = (string: string) => {
-    return string.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
-  };
+  try {
+    const encoder = new TextEncoder();
+    const stream = new TransformStream();
+    const writer = stream.writable.getWriter();
 
-  const createFlexibleHtmlRegex = (searchBlock: string) => {
-    let searchRegex = escapeRegExp(searchBlock)
-      .replace(/\s+/g, '\\s*')
-      .replace(/>\s*</g, '>\\s*<')
-      .replace(/\s*>/g, '\\s*>');
-    
-    return new RegExp(searchRegex, 'g');
-  };
+    const response = new NextResponse(stream.readable, {
+      headers: {
+        "Content-Type": "text/plain; charset=utf-8",
+        "Cache-Control": "no-cache",
+        Connection: "keep-alive",
+      },
+    });
 
-  const selectedProvider = await getBestProvider(selectedModel.value, provider)
+    (async () => {
+      try {
+        const client = new InferenceClient(token);
 
-  try {
-    const systemPrompt = FOLLOW_UP_SYSTEM_PROMPT + (isNew ? PROMPT_FOR_PROJECT_NAME : "");
-    const userContext = "You are modifying the HTML file based on the user's request.";
-    
-    const allPages = pages || [];
-    const pagesContext = allPages
-      .map((p: Page) => `- ${p.path}\n${p.html}`)
-      .join("\n\n");
-    
-    const assistantContext = `${
-      selectedElementHtml
-        ? `\n\nYou have to update ONLY the following element, NOTHING ELSE: \n\n\`\`\`html\n${selectedElementHtml}\n\`\`\` Could be in multiple pages, if so, update all the pages.`
-        : ""
-    }. Current pages (${allPages.length} total): ${pagesContext}. ${files?.length > 0 ? `Available images: ${files?.map((f: string) => f).join(', ')}.` : ""}`;
-    
-    const estimatedInputTokens = estimateInputTokens(systemPrompt, prompt, userContext + assistantContext);
-    const dynamicMaxTokens = calculateMaxTokens(selectedProvider, estimatedInputTokens, false);
-    const providerConfig = getProviderSpecificConfig(selectedProvider, dynamicMaxTokens);
-    
-    const chatCompletion = client.chatCompletionStream(
-      {
-        model: selectedModel.value,
-        provider: selectedProvider.provider,
-        messages: [
-          {
-            role: "system",
-            content: systemPrompt,
-          },
-          {
-            role: "user",
-            content: userContext,
-          },
-          {
-            role: "assistant",
-            content: assistantContext,
-          },
-          {
-            role: "user",
-            content: prompt,
-          },
-        ],
-        ...providerConfig,
-      },
-      billTo ? { billTo } : {}
-    );
+        const systemPrompt = FOLLOW_UP_SYSTEM_PROMPT + (isNew ? PROMPT_FOR_PROJECT_NAME : "");
+        const userContext = "You are modifying the HTML file based on the user's request.";
 
-    // Set up timeout for AI streaming
-    let chunk = "";
-    let timeoutId: NodeJS.Timeout | null = null;
-    let isTimedOut = false;
+        const allPages = pages || [];
+        const pagesContext = allPages
+          .map((p: Page) => `- ${p.path}\n${p.html}`)
+          .join("\n\n");
 
-    const timeoutPromise = new Promise<never>((_, reject) => {
-      timeoutId = setTimeout(() => {
-        isTimedOut = true;
-        reject(new Error('Request timeout: The AI model took too long to respond. Please try again with a simpler prompt or try a different model.'));
-      }, REQUEST_TIMEOUT);
-    });
+        const assistantContext = `${selectedElementHtml
+            ? `\n\nYou have to update ONLY the following element, NOTHING ELSE: \n\n\`\`\`html\n${selectedElementHtml}\n\`\`\` Could be in multiple pages, if so, update all the pages.`
+            : ""
+          }. Current pages (${allPages.length} total): ${pagesContext}. ${files?.length > 0 ? `Available images: ${files?.map((f: string) => f).join(', ')}.` : ""}`;
 
-    try {
-      await Promise.race([
-        (async () => {
-          while (true) {
-            const { done, value } = await chatCompletion.next();
-            if (done) {
-              break;
-            }
-
-            const deltaContent = value.choices[0]?.delta?.content;
-            if (deltaContent) {
-              chunk += deltaContent;
-            }
-          }
-        })(),
-        timeoutPromise
-      ]);
-
-      // Clear timeout if successful
-      if (timeoutId) clearTimeout(timeoutId);
-    } catch (timeoutError: any) {
-      console.error("++TIMEOUT ERROR++", timeoutError);
-      console.error("++TIMEOUT ERROR MESSAGE++", timeoutError.message);
-      // Clear timeout on error
-      if (timeoutId) clearTimeout(timeoutId);
-      
-      if (isTimedOut || timeoutError.message?.includes('timeout') || timeoutError.message?.includes('Request timeout')) {
-        return NextResponse.json(
+        const estimatedInputTokens = estimateInputTokens(systemPrompt, prompt, userContext + assistantContext);
+        const dynamicMaxTokens = calculateMaxTokens(selectedProvider, estimatedInputTokens, false);
+        const providerConfig = getProviderSpecificConfig(selectedProvider, dynamicMaxTokens);
+
+        const chatCompletion = client.chatCompletionStream(
           {
-            ok: false,
-            message: "Request timeout: The AI model took too long to respond. Please try again with a simpler prompt or try a different model.",
+            model: selectedModel.value,
+            provider: selectedProvider.provider,
+            messages: [
+              {
+                role: "system",
+                content: systemPrompt,
+              },
+              {
+                role: "user",
+                content: userContext,
+              },
+              {
+                role: "assistant",
+                content: assistantContext,
+              },
+              {
+                role: "user",
+                content: prompt,
+              },
+            ],
+            ...providerConfig,
           },
-          { status: 504 }
+          billTo ? { billTo } : {}
         );
-      }
-      throw timeoutError;
-    }
-    if (!chunk) {
-      return NextResponse.json(
-        { ok: false, message: "No content returned from the model" },
-        { status: 400 }
-      );
-    }
-
-    if (chunk) {
-      const updatedLines: number[][] = [];
-      let newHtml = "";
-      const updatedPages = [...(pages || [])];
 
-      const updateFileRegex = new RegExp(`${UPDATE_FILE_START.replace(/[.*+?^${}()|[\]\\]/g, '\\$&')}([^\\s]+)\\s*${UPDATE_FILE_END.replace(/[.*+?^${}()|[\]\\]/g, '\\$&')}([\\s\\S]*?)(?=${UPDATE_FILE_START.replace(/[.*+?^${}()|[\]\\]/g, '\\$&')}|${NEW_FILE_START.replace(/[.*+?^${}()|[\]\\]/g, '\\$&')}|$)`, 'g');
-      let updateFileMatch;
-      
-      while ((updateFileMatch = updateFileRegex.exec(chunk)) !== null) {
-        const [, filePath, fileContent] = updateFileMatch;
-        
-        const pageIndex = updatedPages.findIndex(p => p.path === filePath);
-        if (pageIndex !== -1) {
-          let pageHtml = updatedPages[pageIndex].html;
-          
-          let processedContent = fileContent;
-          const htmlMatch = fileContent.match(/```html\s*([\s\S]*?)\s*```/);
-          if (htmlMatch) {
-            processedContent = htmlMatch[1];
-          }
-          let position = 0;
-          let moreBlocks = true;
-
-          while (moreBlocks) {
-            const searchStartIndex = processedContent.indexOf(SEARCH_START, position);
-            if (searchStartIndex === -1) {
-              moreBlocks = false;
-              continue;
-            }
-
-            const dividerIndex = processedContent.indexOf(DIVIDER, searchStartIndex);
-            if (dividerIndex === -1) {
-              moreBlocks = false;
-              continue;
-            }
-
-            const replaceEndIndex = processedContent.indexOf(REPLACE_END, dividerIndex);
-            if (replaceEndIndex === -1) {
-              moreBlocks = false;
-              continue;
-            }
-
-            const searchBlock = processedContent.substring(
-              searchStartIndex + SEARCH_START.length,
-              dividerIndex
-            );
-            const replaceBlock = processedContent.substring(
-              dividerIndex + DIVIDER.length,
-              replaceEndIndex
-            );
-
-            if (searchBlock.trim() === "") {
-              pageHtml = `${replaceBlock}\n${pageHtml}`;
-              updatedLines.push([1, replaceBlock.split("\n").length]);
-            } else {
-              const regex = createFlexibleHtmlRegex(searchBlock);
-              const match = regex.exec(pageHtml);
-              
-              if (match) {
-                const matchedText = match[0];
-                const beforeText = pageHtml.substring(0, match.index);
-                const startLineNumber = beforeText.split("\n").length;
-                const replaceLines = replaceBlock.split("\n").length;
-                const endLineNumber = startLineNumber + replaceLines - 1;
-
-                updatedLines.push([startLineNumber, endLineNumber]);
-                pageHtml = pageHtml.replace(matchedText, replaceBlock);
-              }
-            }
-
-            position = replaceEndIndex + REPLACE_END.length;
+        // Stream the response chunks to the client
+        while (true) {
+          const { done, value } = await chatCompletion.next();
+          if (done) {
+            break;
           }
 
-          updatedPages[pageIndex].html = pageHtml;
-          
-          if (filePath === '/' || filePath === '/index' || filePath === 'index' || filePath === 'index.html') {
-            newHtml = pageHtml;
+          const chunk = value.choices[0]?.delta?.content;
+          if (chunk) {
+            await writer.write(encoder.encode(chunk));
           }
         }
-      }
-
-      const newFileRegex = new RegExp(`${NEW_FILE_START.replace(/[.*+?^${}()|[\]\\]/g, '\\$&')}([^\\s]+)\\s*${NEW_FILE_END.replace(/[.*+?^${}()|[\]\\]/g, '\\$&')}([\\s\\S]*?)(?=${UPDATE_FILE_START.replace(/[.*+?^${}()|[\]\\]/g, '\\$&')}|${NEW_FILE_START.replace(/[.*+?^${}()|[\]\\]/g, '\\$&')}|$)`, 'g');
-      let newFileMatch;
-      
-      while ((newFileMatch = newFileRegex.exec(chunk)) !== null) {
-        const [, filePath, fileContent] = newFileMatch;
-        
-        let fileData = fileContent;
-        // Try to extract content from code blocks
-        const htmlMatch = fileContent.match(/```html\s*([\s\S]*?)\s*```/);
-        const cssMatch = fileContent.match(/```css\s*([\s\S]*?)\s*```/);
-        const jsMatch = fileContent.match(/```javascript\s*([\s\S]*?)\s*```/);
-        
-        if (htmlMatch) {
-          fileData = htmlMatch[1];
-        } else if (cssMatch) {
-          fileData = cssMatch[1];
-        } else if (jsMatch) {
-          fileData = jsMatch[1];
-        }
-        
-        const existingFileIndex = updatedPages.findIndex(p => p.path === filePath);
-        
-        if (existingFileIndex !== -1) {
-          updatedPages[existingFileIndex] = {
-            path: filePath,
-            html: fileData.trim()
-          };
-        } else {
-          updatedPages.push({
-            path: filePath,
-            html: fileData.trim()
-          });
-        }
-      }
-
-      if (updatedPages.length === pages?.length && !chunk.includes(UPDATE_FILE_START)) {
-        let position = 0;
-        let moreBlocks = true;
-
-        while (moreBlocks) {
-          const searchStartIndex = chunk.indexOf(SEARCH_START, position);
-          if (searchStartIndex === -1) {
-            moreBlocks = false;
-            continue;
-          }
 
-          const dividerIndex = chunk.indexOf(DIVIDER, searchStartIndex);
-          if (dividerIndex === -1) {
-            moreBlocks = false;
-            continue;
-          }
+        await writer.write(encoder.encode(`\n___METADATA_START___\n${JSON.stringify({
+          repoId,
+          isNew,
+          userName: user.name,
+        })}\n___METADATA_END___\n`));
 
-          const replaceEndIndex = chunk.indexOf(REPLACE_END, dividerIndex);
-          if (replaceEndIndex === -1) {
-            moreBlocks = false;
-            continue;
-          }
-
-          const searchBlock = chunk.substring(
-            searchStartIndex + SEARCH_START.length,
-            dividerIndex
+        await writer.close();
+      } catch (error: any) {
+        if (error.message?.includes("exceeded your monthly included credits")) {
+          await writer.write(
+            encoder.encode(
+              JSON.stringify({
+                ok: false,
+                openProModal: true,
+                message: error.message,
+              })
+            )
           );
-          const replaceBlock = chunk.substring(
-            dividerIndex + DIVIDER.length,
-            replaceEndIndex
+        } else if (error?.message?.includes("inference provider information")) {
+          await writer.write(
+            encoder.encode(
+              JSON.stringify({
+                ok: false,
+                openSelectProvider: true,
+                message: error.message,
+              })
+            )
+          );
+        } else {
+          await writer.write(
+            encoder.encode(
+              JSON.stringify({
+                ok: false,
+                message:
+                  error.message ||
+                  "An error occurred while processing your request.",
+              })
+            )
           );
-
-          if (searchBlock.trim() === "") {
-            newHtml = `${replaceBlock}\n${newHtml}`;
-            updatedLines.push([1, replaceBlock.split("\n").length]);
-          } else {
-            const regex = createFlexibleHtmlRegex(searchBlock);
-            const match = regex.exec(newHtml);
-            
-            if (match) {
-              const matchedText = match[0];
-              const beforeText = newHtml.substring(0, match.index);
-              const startLineNumber = beforeText.split("\n").length;
-              const replaceLines = replaceBlock.split("\n").length;
-              const endLineNumber = startLineNumber + replaceLines - 1;
-
-              updatedLines.push([startLineNumber, endLineNumber]);
-              newHtml = newHtml.replace(matchedText, replaceBlock);
-            }
-          }
-
-          position = replaceEndIndex + REPLACE_END.length;
-        }
-
-        const mainPageIndex = updatedPages.findIndex(p => p.path === '/' || p.path === '/index' || p.path === 'index');
-        if (mainPageIndex !== -1) {
-          updatedPages[mainPageIndex].html = newHtml;
-        }
-      }
-
-      const files: File[] = [];
-      updatedPages.forEach((page: Page) => {
-        let mimeType = "text/html";
-        if (page.path.endsWith(".css")) {
-          mimeType = "text/css";
-        } else if (page.path.endsWith(".js")) {
-          mimeType = "text/javascript";
-        } else if (page.path.endsWith(".json")) {
-          mimeType = "application/json";
         }
-        const content = (mimeType === "text/html" && isIndexPage(page.path)) && isNew
-          ? injectDeepSiteBadge(page.html) 
-          : page.html;
-        const file = new File([content], page.path, { type: mimeType });
-        files.push(file);
-      });
-
-      if (isNew) {
-        const projectName = chunk.match(/<<<<<<< PROJECT_NAME_START\s*([\s\S]*?)\s*>>>>>>> PROJECT_NAME_END/)?.[1]?.trim();
-        const formattedTitle = projectName?.toLowerCase()
-          .replace(/[^a-z0-9]+/g, "-")
-          .split("-")
-          .filter(Boolean)
-          .join("-")
-          .slice(0, 96);
-        const repo: RepoDesignation = {
-          type: "space",
-          name: `${user.name}/${formattedTitle}`,
-        };
-        
+      } finally {
         try {
-          const { repoUrl} = await createRepo({
-            repo,
-            accessToken: user.token as string,
-          });
-          repoId = repoUrl.split("/").slice(-2).join("/");
-        } catch (createRepoError: any) {
-          console.error("++CREATE REPO ERROR++", createRepoError);
-          throw new Error(`Failed to create repository: ${createRepoError.message || 'Unknown error'}`);
+          await writer?.close();
+        } catch {
+          // ignore
         }
-        
-        const colorFrom = COLORS[Math.floor(Math.random() * COLORS.length)];
-        const colorTo = COLORS[Math.floor(Math.random() * COLORS.length)];
-        const README = `---
-title: ${projectName}
-colorFrom: ${colorFrom}
-colorTo: ${colorTo}
-emoji: 🐳
-sdk: static
-pinned: false
-tags:
-  - deepsite-v3
----
-
-# Welcome to your new DeepSite project!
-This project was created with [DeepSite](https://huggingface.co/deepsite).
-      `;
-        files.push(new File([README], "README.md", { type: "text/markdown" }));
-      }
-
-      let response;
-      try {
-        // Add a timeout wrapper for the upload
-        const uploadPromise = uploadFiles({
-          repo: {
-            type: "space",
-            name: repoId,
-          },
-          files,
-          commitTitle: prompt,
-          accessToken: user.token as string,
-        });
-        
-        const uploadTimeout = new Promise<never>((_, reject) => {
-          setTimeout(() => {
-            reject(new Error('Upload operation timed out'));
-          }, 180000); // 3 minutes timeout for upload
-        });
-        
-        response = await Promise.race([uploadPromise, uploadTimeout]);
-      } catch (uploadError: any) {
-        console.error("++UPLOAD FILES ERROR++", uploadError);
-        console.error("++UPLOAD FILES ERROR MESSAGE++", uploadError.message);
-        console.error("++REPO ID++", repoId);
-        
-        // If it's a timeout, files might have been uploaded but we didn't get response
-        if (uploadError.message?.includes('timed out') || uploadError.message?.includes('timeout')) {
-          console.warn("++UPLOAD TIMEOUT - Files may have been uploaded++");
-          // Return a partial success response
-          return NextResponse.json({
-            ok: true,
-            updatedLines,
-            pages: updatedPages,
-            repoId,
-            commit: {
-              title: prompt,
-              oid: 'timeout',
-              timedOut: true,
-            }
-          });
-        }
-        
-        throw new Error(`Failed to upload files to repository: ${uploadError.message || 'Unknown error'}`);
-      }
-      const responseData: any = {
-        ok: true,
-        updatedLines,
-        pages: updatedPages,
-        repoId,
-      };
-
-      if (response && response.commit) {
-        responseData.commit = {
-          ...response.commit,
-          title: prompt,
-        };
-      } else {
-        responseData.commit = {
-          title: prompt,
-          oid: 'unknown',
-        };
       }
+    })();
 
-      return NextResponse.json(responseData);
-    } else {
-      return NextResponse.json(
-        { ok: false, message: "No content returned from the model" },
-        { status: 400 }
-      );
-    }
+    return response;
   } catch (error: any) {
-    console.error("++ERROR++", error);
-    console.error("++ERROR MESSAGE++", error.message);
-    if (error.message?.includes('timeout') || error.message?.includes('Request timeout')) {
-      return NextResponse.json(
-        {
-          ok: false,
-          message: "Request timeout: The operation took too long to complete. Please try again with a simpler request or try a different model.",
-        },
-        { status: 504 }
-      );
-    }
-    if (error.message?.includes("exceeded your monthly included credits")) {
-      return NextResponse.json(
-        {
-          ok: false,
-          openProModal: true,
-          message: error.message,
-        },
-        { status: 402 }
-      );
-    }
     return NextResponse.json(
       {
         ok: false,

app/api/me/projects/[namespace]/[repoId]/update/route.ts

@@ -0,0 +1,141 @@
+import { NextRequest, NextResponse } from "next/server";
+import { createRepo, RepoDesignation, uploadFiles } from "@huggingface/hub";
+
+import { isAuthenticated } from "@/lib/auth";
+import { Page } from "@/types";
+import { COLORS } from "@/lib/utils";
+import { injectDeepSiteBadge, isIndexPage } from "@/lib/inject-badge";
+import { pagesToFiles } from "@/lib/format-ai-response";
+
+/**
+ * UPDATE route - for updating existing projects or creating new ones after AI streaming
+ * This route handles the HuggingFace upload after client-side AI response processing
+ */
+export async function PUT(
+  req: NextRequest,
+  { params }: { params: Promise<{ namespace: string; repoId: string }> }
+) {
+  const user = await isAuthenticated();
+  if (user instanceof NextResponse || !user) {
+    return NextResponse.json({ message: "Unauthorized" }, { status: 401 });
+  }
+
+  const param = await params;
+  let { namespace, repoId } = param;
+  const { pages, commitTitle = "AI-generated changes", isNew, projectName } = await req.json();
+
+  if (!pages || !Array.isArray(pages) || pages.length === 0) {
+    return NextResponse.json(
+      { ok: false, error: "Pages are required" },
+      { status: 400 }
+    );
+  }
+
+  try {
+    let files: File[];
+    
+    if (isNew) {
+      // Creating a new project
+      const title = projectName || "DeepSite Project";
+      const formattedTitle = title
+        .toLowerCase()
+        .replace(/[^a-z0-9]+/g, "-")
+        .split("-")
+        .filter(Boolean)
+        .join("-")
+        .slice(0, 96);
+
+      const repo: RepoDesignation = {
+        type: "space",
+        name: `${user.name}/${formattedTitle}`,
+      };
+
+      try {
+        const { repoUrl } = await createRepo({
+          repo,
+          accessToken: user.token as string,
+        });
+        namespace = user.name;
+        repoId = repoUrl.split("/").slice(-2).join("/").split("/")[1];
+      } catch (createRepoError: any) {
+        return NextResponse.json(
+          {
+            ok: false,
+            error: `Failed to create repository: ${createRepoError.message || 'Unknown error'}`,
+          },
+          { status: 500 }
+        );
+      }
+
+      // Prepare files with badge injection for new projects
+      files = [];
+      pages.forEach((page: Page) => {
+        let mimeType = "text/html";
+        if (page.path.endsWith(".css")) {
+          mimeType = "text/css";
+        } else if (page.path.endsWith(".js")) {
+          mimeType = "text/javascript";
+        } else if (page.path.endsWith(".json")) {
+          mimeType = "application/json";
+        }
+        const content = (mimeType === "text/html" && isIndexPage(page.path))
+          ? injectDeepSiteBadge(page.html)
+          : page.html;
+        const file = new File([content], page.path, { type: mimeType });
+        files.push(file);
+      });
+
+      // Add README.md for new projects
+      const colorFrom = COLORS[Math.floor(Math.random() * COLORS.length)];
+      const colorTo = COLORS[Math.floor(Math.random() * COLORS.length)];
+      const README = `---
+title: ${title}
+colorFrom: ${colorFrom}
+colorTo: ${colorTo}
+emoji: 🐳
+sdk: static
+pinned: false
+tags:
+  - deepsite-v3
+---
+
+# Welcome to your new DeepSite project!
+This project was created with [DeepSite](https://huggingface.co/deepsite).
+`;
+      files.push(new File([README], "README.md", { type: "text/markdown" }));
+    } else {
+      // Updating existing project - no badge injection
+      files = pagesToFiles(pages);
+    }
+
+    const response = await uploadFiles({
+      repo: {
+        type: "space",
+        name: `${namespace}/${repoId}`,
+      },
+      files,
+      commitTitle,
+      accessToken: user.token as string,
+    });
+
+    return NextResponse.json({
+      ok: true,
+      pages,
+      repoId: `${namespace}/${repoId}`,
+      commit: {
+        ...response.commit,
+        title: commitTitle,
+      }
+    });
+  } catch (error: any) {
+    console.error("Error updating project:", error);
+    return NextResponse.json(
+      {
+        ok: false,
+        error: error.message || "Failed to update project",
+      },
+      { status: 500 }
+    );
+  }
+}
+

components/editor/ask-ai/index.tsx

@@ -100,7 +100,8 @@ export const AskAi = ({
         prompt,
         enhancedSettings,
         redesignMarkdown,
-        !!user
+        !!user,
+        user?.name
       );
 
       if (result?.error) {

components/editor/preview/index.tsx

@@ -72,7 +72,6 @@ export const Preview = ({ isNew }: { isNew: boolean }) => {
     (html: string): string => {
       if (!html) return html;
 
-      // Find all CSS and JS files (including those in subdirectories)
       const cssFiles = pages.filter(
         (p) => p.path.endsWith(".css") && p.path !== previewPageData?.path
       );
@@ -102,11 +101,9 @@ export const Preview = ({ isNew }: { isNew: boolean }) => {
             `<head>\n${allCssContent}`
           );
         } else {
-          // If no head tag, prepend to document
           modifiedHtml = allCssContent + "\n" + modifiedHtml;
         }
 
-        // Remove all link tags that reference CSS files we're injecting
         cssFiles.forEach((file) => {
           const escapedPath = file.path.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");
           modifiedHtml = modifiedHtml.replace(
@@ -119,7 +116,6 @@ export const Preview = ({ isNew }: { isNew: boolean }) => {
         });
       }
 
-      // Inject all JS files
       if (jsFiles.length > 0) {
         const allJsContent = jsFiles
           .map(
@@ -136,11 +132,9 @@ export const Preview = ({ isNew }: { isNew: boolean }) => {
         } else if (modifiedHtml.includes("<body>")) {
           modifiedHtml = modifiedHtml + allJsContent;
         } else {
-          // If no body tag, append to document
           modifiedHtml = modifiedHtml + "\n" + allJsContent;
         }
 
-        // Remove all script tags that reference JS files we're injecting
         jsFiles.forEach((file) => {
           const escapedPath = file.path.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");
           modifiedHtml = modifiedHtml.replace(
@@ -214,7 +208,6 @@ export const Preview = ({ isNew }: { isNew: boolean }) => {
     if (iframeRef?.current?.contentDocument) {
       const iframeDocument = iframeRef.current.contentDocument;
 
-      // Use event delegation to catch clicks on anchors in both light and shadow DOM
       iframeDocument.addEventListener(
         "click",
         handleCustomNavigation as any,
@@ -323,18 +316,11 @@ export const Preview = ({ isNew }: { isNew: boolean }) => {
         const path = event.composedPath();
         const targetElement = path[0] as HTMLElement;
 
-        console.log(
-          "[handleClick] Target element:",
-          targetElement.tagName,
-          targetElement
-        );
-
         const findClosestAnchor = (
           element: HTMLElement
         ): HTMLAnchorElement | null => {
           let current: HTMLElement | null = element;
           while (current) {
-            console.log("[handleClick] Checking element:", current.tagName);
             if (current.tagName?.toUpperCase() === "A") {
               return current as HTMLAnchorElement;
             }
@@ -342,13 +328,9 @@ export const Preview = ({ isNew }: { isNew: boolean }) => {
               break;
             }
             const parent: Node | null = current.parentNode;
-            // Use nodeType to check - works across iframe boundaries
-            // nodeType 1 = Element, nodeType 11 = DocumentFragment (including ShadowRoot)
             if (parent && parent.nodeType === 11) {
-              // ShadowRoot
               current = (parent as ShadowRoot).host as HTMLElement;
             } else if (parent && parent.nodeType === 1) {
-              // Element node
               current = parent as HTMLElement;
             } else {
               break;
@@ -359,8 +341,6 @@ export const Preview = ({ isNew }: { isNew: boolean }) => {
 
         const anchorElement = findClosestAnchor(targetElement);
 
-        console.log("[handleClick] Found anchor:", anchorElement);
-
         if (anchorElement) {
           return;
         }
@@ -379,49 +359,23 @@ export const Preview = ({ isNew }: { isNew: boolean }) => {
         const path = event.composedPath();
         const actualTarget = path[0] as HTMLElement;
 
-        console.log(
-          "[handleCustomNavigation] Click detected in iframe:",
-          actualTarget.tagName,
-          actualTarget
-        );
-
         const findClosestAnchor = (
           element: HTMLElement
         ): HTMLAnchorElement | null => {
           let current: HTMLElement | null = element;
           while (current) {
-            console.log(
-              "[handleCustomNavigation] Checking element:",
-              current.tagName,
-              current
-            );
             if (current.tagName?.toUpperCase() === "A") {
-              console.log("[handleCustomNavigation] Found anchor!", current);
               return current as HTMLAnchorElement;
             }
             if (current === iframeDocument.body) {
-              console.log("[handleCustomNavigation] Reached body, stopping");
               break;
             }
             const parent: Node | null = current.parentNode;
-            console.log(
-              "[handleCustomNavigation] Parent node:",
-              parent,
-              "nodeType:",
-              parent?.nodeType
-            );
-            // Use nodeType to check - works across iframe boundaries
-            // nodeType 1 = Element, nodeType 11 = DocumentFragment (including ShadowRoot)
             if (parent && parent.nodeType === 11) {
-              // ShadowRoot
               current = (parent as ShadowRoot).host as HTMLElement;
             } else if (parent && parent.nodeType === 1) {
-              // Element node
               current = parent as HTMLElement;
             } else {
-              console.log(
-                "[handleCustomNavigation] Parent is not an element node, breaking"
-              );
               break;
             }
           }
@@ -429,7 +383,6 @@ export const Preview = ({ isNew }: { isNew: boolean }) => {
         };
 
         const anchorElement = findClosestAnchor(actualTarget);
-        console.log("[handleCustomNavigation] Anchor element:", anchorElement);
         if (anchorElement) {
           let href = anchorElement.getAttribute("href");
           if (href) {
@@ -573,7 +526,6 @@ export const Preview = ({ isNew }: { isNew: boolean }) => {
                         }
                       );
                     }
-                    // Set up event listeners after iframe loads
                     setupIframeListeners();
                   }
                 : undefined
@@ -583,7 +535,7 @@ export const Preview = ({ isNew }: { isNew: boolean }) => {
           />
           {!isNew && (
             <>
-              <div
+              {/* <div
                 className={classNames(
                   "w-full h-full flex items-center justify-center absolute left-0 top-0 bg-black/40 backdrop-blur-lg transition-all duration-200",
                   {
@@ -601,7 +553,7 @@ export const Preview = ({ isNew }: { isNew: boolean }) => {
                   <AnimatedBlobs />
                   <AnimatedBlobs />
                 </div>
-              </div>
+              </div> */}
               <HistoryNotification
                 isVisible={!!currentCommit}
                 isPromotingVersion={isPromotingVersion}

hooks/useAi.ts

@@ -124,27 +124,37 @@ export const useAi = (onScrollToBottom?: () => void) => {
     client.setQueryData(["ai.model"], newModel);
   };
 
-  const createNewProject = async (prompt: string, htmlPages: Page[], projectName: string | undefined, isLoggedIn?: boolean) => {
-    if (isLoggedIn) {
-      api.post("/me/projects", {
-        title: projectName,
-        pages: htmlPages,
-        prompt,
-      })
-      .then((response) => {
-        if (response.data.ok) {
-          setIsAiWorking(false);
-          router.replace(`/${response.data.space.project.space_id}`);
-          setProject(response.data.space);
-          setProjects([...projects, response.data.space]);
-          toast.success("AI responded successfully");
-          if (audio.current) audio.current.play();
+  const createNewProject = async (prompt: string, htmlPages: Page[], projectName: string | undefined, isLoggedIn?: boolean, userName?: string) => {
+    if (isLoggedIn && userName) {
+      try {
+        const uploadRequest = await fetch(`/deepsite/api/me/projects/${userName}/new/update`, {
+          method: "PUT",
+          body: JSON.stringify({
+            pages: htmlPages,
+            commitTitle: prompt,
+            isNew: true,
+            projectName,
+          }),
+          headers: {
+            "Content-Type": "application/json",
+            "Authorization": `Bearer ${token}`,
+          },
+        });
+
+        const uploadRes = await uploadRequest.json();
+        
+        if (!uploadRequest.ok || !uploadRes.ok) {
+          throw new Error(uploadRes.error || "Failed to create project");
         }
-      })
-      .catch((error) => {
+
         setIsAiWorking(false);
-        toast.error(error?.response?.data?.message || error?.message || "Failed to create project");
-      });
+        router.replace(`/${uploadRes.repoId}`);
+        toast.success("AI responded successfully");
+        if (audio.current) audio.current.play();
+      } catch (error: any) {
+        setIsAiWorking(false);
+        toast.error(error?.message || "Failed to create project");
+      }
     } else {
       setIsAiWorking(false);
       toast.success("AI responded successfully");
@@ -152,7 +162,7 @@ export const useAi = (onScrollToBottom?: () => void) => {
     }
   }
   
-  const callAiNewProject = async (prompt: string, enhancedSettings?: EnhancedSettings, redesignMarkdown?: string, isLoggedIn?: boolean) => {
+  const callAiNewProject = async (prompt: string, enhancedSettings?: EnhancedSettings, redesignMarkdown?: string, isLoggedIn?: boolean, userName?: string) => {
     if (isAiWorking) return;
     if (!redesignMarkdown && !prompt.trim()) return;
     
@@ -218,7 +228,7 @@ export const useAi = (onScrollToBottom?: () => void) => {
             setPages(newPages);
             setLastSavedPages([...newPages]);
             if (newPages.length > 0 && !isTheSameHtml(newPages[0].html)) {
-              createNewProject(prompt, newPages, projectName, isLoggedIn);
+              createNewProject(prompt, newPages, projectName, isLoggedIn, userName);
             }
             setPrompts([...prompts, prompt]);
 
@@ -311,118 +321,162 @@ export const useAi = (onScrollToBottom?: () => void) => {
       });
 
       if (request && request.body) {
-        if (request.status === 504) {
-          console.warn("++504 GATEWAY TIMEOUT - Upload likely succeeded but response timed out++");
-          console.warn("++REQUEST STATUS++", request.status);
-          console.warn("++REQUEST BODY++", request.body);
-          
-          setIsAiWorking(false);
-          
-          if (isNew) {
-            toast.error("The request timed out. Your project may have been created. Please check your HuggingFace spaces.");
-            return { error: "gateway_timeout", message: "Request timed out after upload" };
-          }
+        const reader = request.body.getReader();
+        const decoder = new TextDecoder("utf-8");
+        let contentResponse = "";
+        let metadata: any = null;
+
+        const read = async (): Promise<any> => {
+          const { done, value } = await reader.read();
           
-          toast.success("Changes saved! Refreshing page to show updates...", { duration: 3000 });
-          setTimeout(() => {
-            window.location.reload();
-          }, 1000);
-          return { success: true, timedOut: true };
-        }
-        
-        const clonedRequest = request.clone();
-        let res;
-        try {
-          res = await request.json();
-        } catch (jsonError: any) {
-          try {
-            const text = await clonedRequest.text();
+          if (done) {
+            const metadataMatch = contentResponse.match(/___METADATA_START___([\s\S]*?)___METADATA_END___/);
+            if (metadataMatch) {
+              try {
+                metadata = JSON.parse(metadataMatch[1]);
+                contentResponse = contentResponse.replace(/___METADATA_START___[\s\S]*?___METADATA_END___/, '').trim();
+              } catch (e) {
+                console.error("Failed to parse metadata", e);
+              }
+            }
+
+            const trimmedResponse = contentResponse.trim();
+            if (trimmedResponse.startsWith("{") && trimmedResponse.endsWith("}")) {
+              try {
+                const jsonResponse = JSON.parse(trimmedResponse);
+                if (jsonResponse && !jsonResponse.ok) {
+                  setIsAiWorking(false);
+                  if (jsonResponse.openLogin) {
+                    return { error: "login_required" };
+                  } else if (jsonResponse.openSelectProvider) {
+                    return { error: "provider_required", message: jsonResponse.message };
+                  } else if (jsonResponse.openProModal) {
+                    return { error: "pro_required" };
+                  } else {
+                    toast.error(jsonResponse.message);
+                    return { error: "api_error", message: jsonResponse.message };
+                  }
+                }
+              } catch (e) {
+                // Not JSON, continue with normal processing
+              }
+            }
             
-            // Check if it's a CloudFront/gateway timeout in the HTML
-            if (text.includes("504") || text.includes("Gateway Timeout") || text.includes("gateway timeout")) {
-              console.warn("++DETECTED 504 IN HTML RESPONSE++");
-              setIsAiWorking(false);
-              
-              if (isNew) {
-                toast.error("The request timed out. Your project may have been created. Please check your HuggingFace spaces.");
-                return { error: "gateway_timeout", message: "Request timed out after upload" };
+            const { processAiResponse, extractProjectName } = await import("@/lib/format-ai-response");
+            const { updatedPages, updatedLines } = processAiResponse(contentResponse, pagesToSend);
+            
+            const updatedPagesMap = new Map(updatedPages.map((p: Page) => [p.path, p]));
+            const mergedPages: Page[] = pages.map(page => 
+              updatedPagesMap.has(page.path) ? updatedPagesMap.get(page.path)! : page
+            );
+            updatedPages.forEach((page: Page) => {
+              if (!pages.find(p => p.path === page.path)) {
+                mergedPages.push(page);
+              }
+            });
+
+            let projectName = null;
+            if (isNew) {
+              projectName = extractProjectName(contentResponse);
+              if (!projectName) {
+                projectName = prompt.substring(0, 40).replace(/[^a-zA-Z0-9]/g, "-").slice(0, 40) + "-" + Math.random().toString(36).substring(2, 15);
               }
+            }
+
+            try {
+              const uploadRequest = await fetch(`/deepsite/api/me/projects/${metadata?.userName || 'unknown'}/${isNew ? 'new' : (project?.space_id?.split('/')[1] || 'unknown')}/update`, {
+                method: "PUT",
+                body: JSON.stringify({
+                  pages: mergedPages,
+                  commitTitle: prompt,
+                  isNew,
+                  projectName,
+                }),
+                headers: {
+                  "Content-Type": "application/json",
+                  "Authorization": `Bearer ${token}`,
+                },
+              });
+
+              const uploadRes = await uploadRequest.json();
               
-              toast.success("Changes saved! Refreshing page to show updates...", { duration: 3000 });
-              setTimeout(() => {
-                window.location.reload();
-              }, 1000);
-              return { success: true, timedOut: true };
+              if (!uploadRequest.ok || !uploadRes.ok) {
+                throw new Error(uploadRes.error || "Failed to upload to HuggingFace");
+              }
+
+              toast.success("AI responded successfully");
+              const iframe = document.getElementById("preview-iframe") as HTMLIFrameElement;
+
+              if (isNew && uploadRes.repoId) {
+                router.push(`/${uploadRes.repoId}`);
+                setIsAiWorking(false);
+              } else {
+                setPages(mergedPages);
+                setLastSavedPages([...mergedPages]);
+                setCommits([uploadRes.commit, ...commits]);
+                setPrompts([...prompts, prompt]);
+                setSelectedElement(null);
+                setSelectedFiles([]);
+                setIsEditableModeEnabled(false);
+                setIsAiWorking(false);
+              }
+
+              if (audio.current) audio.current.play();
+              if (iframe) {
+                setTimeout(() => {
+                  iframe.src = iframe.src;
+                }, 500);
+              }
+
+              return { success: true, updatedLines };
+            } catch (uploadError: any) {
+              setIsAiWorking(false);
+              toast.error(uploadError.message || "Failed to save changes");
+              return { error: "upload_error", message: uploadError.message };
             }
-          } catch (textError) {
-          }
-          setIsAiWorking(false);
-          toast.error("Server returned invalid response. Check console for details.");
-          return { error: "invalid_response", message: "Server returned non-JSON response" };
-        }
-        
-        if (!request.ok) {
-          if (res.openLogin) {
-            setIsAiWorking(false);
-            return { error: "login_required" };
-          } else if (res.openSelectProvider) {
-            setIsAiWorking(false);
-            return { error: "provider_required", message: res.message };
-          } else if (res.openProModal) {
-            setIsAiWorking(false);
-            return { error: "pro_required" };
-          } else {
-            setIsAiWorking(false);
-            return { error: "api_error", message: res.message };
           }
-        }
 
-        toast.success("AI responded successfully");
-        const iframe = document.getElementById(
-          "preview-iframe"
-        ) as HTMLIFrameElement;
-
-        if (isNew && res.repoId) {
-          router.push(`/${res.repoId}`);
-          setIsAiWorking(false);
-        } else {
-          const returnedPages = res.pages as Page[];
-          const updatedPagesMap = new Map(returnedPages.map((p: Page) => [p.path, p]));
-          const mergedPages: Page[] = pages.map(page => 
-            updatedPagesMap.has(page.path) ? updatedPagesMap.get(page.path)! : page
-          );
-          returnedPages.forEach((page: Page) => {
-            if (!pages.find(p => p.path === page.path)) {
-              mergedPages.push(page);
+          const chunk = decoder.decode(value, { stream: true });
+          contentResponse += chunk;
+          
+          // Check for error responses during streaming
+          const trimmedResponse = contentResponse.trim();
+          if (trimmedResponse.startsWith("{") && trimmedResponse.endsWith("}")) {
+            try {
+              const jsonResponse = JSON.parse(trimmedResponse);
+              if (jsonResponse && !jsonResponse.ok) {
+                setIsAiWorking(false);
+                if (jsonResponse.openLogin) {
+                  return { error: "login_required" };
+                } else if (jsonResponse.openSelectProvider) {
+                  return { error: "provider_required", message: jsonResponse.message };
+                } else if (jsonResponse.openProModal) {
+                  return { error: "pro_required" };
+                } else {
+                  toast.error(jsonResponse.message);
+                  return { error: "api_error", message: jsonResponse.message };
+                }
+              }
+            } catch (e) {
+              // Not complete JSON yet, continue
             }
-          });
+          }
           
-          setPages(mergedPages);
-          setLastSavedPages([...mergedPages]);
-          setCommits([res.commit, ...commits]);
-          setPrompts(
-            [...prompts, prompt]
-          )
-          setSelectedElement(null);
-          setSelectedFiles([]);
-          // setContextFile(null); not needed yet, keep context for the next request.
-          setIsEditableModeEnabled(false);
-          setIsAiWorking(false);
-        }
-
-        if (audio.current) audio.current.play();
-        if (iframe) {
-          setTimeout(() => {
-            iframe.src = iframe.src;
-          }, 500);
-        }
+          return read();
+        };
 
-        return { success: true, html: res.html, updatedLines: res.updatedLines };
+        return await read();
       }
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
     } catch (error: any) {
       setIsAiWorking(false);
-      toast.error(error.message);
+      setIsThinking(false);
+      setController(null);
+      
+      if (!abortController.signal.aborted) {
+        toast.error(error.message || "Network error occurred");
+      }
+      
       if (error.openLogin) {
         return { error: "login_required" };
       }
@@ -470,7 +524,6 @@ export const useAi = (onScrollToBottom?: () => void) => {
     if (pages.length > 0) {
       setPages(pages);
       if (isStreaming) {
-        // Find new pages that haven't been shown yet (HTML, CSS, JS, etc.)
         const newPages = pages.filter(p => 
           !streamingPagesRef.current.has(p.path)
         );
@@ -480,7 +533,6 @@ export const useAi = (onScrollToBottom?: () => void) => {
           setCurrentPage(newPage.path);
           streamingPagesRef.current.add(newPage.path);
           
-          // Update preview if it's an HTML file not in components folder
           if (newPage.path.endsWith('.html') && !newPage.path.includes('/components/')) {
             setPreviewPage(newPage.path);
           }
@@ -500,41 +552,30 @@ export const useAi = (onScrollToBottom?: () => void) => {
   const extractFileContent = (chunk: string, filePath: string): string => {
     if (!chunk) return "";
     
-    // Remove backticks first
     let content = chunk.trim();
     
-    // Handle different file types
     if (filePath.endsWith('.css')) {
-      // Try to extract CSS from complete code blocks first
       const cssMatch = content.match(/```css\s*([\s\S]*?)\s*```/);
       if (cssMatch) {
         content = cssMatch[1];
       } else {
-        // Handle incomplete code blocks during streaming (remove opening fence)
         content = content.replace(/^```css\s*/i, "");
       }
-      // Remove any remaining backticks
       return content.replace(/```/g, "").trim();
     } else if (filePath.endsWith('.js')) {
-      // Try to extract JavaScript from complete code blocks first
       const jsMatch = content.match(/```(?:javascript|js)\s*([\s\S]*?)\s*```/);
       if (jsMatch) {
         content = jsMatch[1];
       } else {
-        // Handle incomplete code blocks during streaming (remove opening fence)
         content = content.replace(/^```(?:javascript|js)\s*/i, "");
       }
-      // Remove any remaining backticks
       return content.replace(/```/g, "").trim();
     } else {
-      // Handle HTML files
       const htmlMatch = content.match(/```html\s*([\s\S]*?)\s*```/);
       if (htmlMatch) {
         content = htmlMatch[1];
       } else {
-        // Handle incomplete code blocks during streaming (remove opening fence)
         content = content.replace(/^```html\s*/i, "");
-        // Try to find HTML starting with DOCTYPE
         const doctypeMatch = content.match(/<!DOCTYPE html>[\s\S]*/);
         if (doctypeMatch) {
           content = doctypeMatch[0];

lib/format-ai-response.ts

@@ -0,0 +1,255 @@
+/* eslint-disable @typescript-eslint/no-explicit-any */
+import { Page } from "@/types";
+import {
+  DIVIDER,
+  NEW_FILE_END,
+  NEW_FILE_START,
+  REPLACE_END,
+  SEARCH_START,
+  UPDATE_FILE_END,
+  UPDATE_FILE_START,
+} from "./prompts";
+
+/**
+ * Escape special regex characters in a string
+ */
+export const escapeRegExp = (string: string): string => {
+  return string.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
+};
+
+/**
+ * Create a flexible HTML regex that accounts for whitespace variations
+ */
+export const createFlexibleHtmlRegex = (searchBlock: string): RegExp => {
+  let searchRegex = escapeRegExp(searchBlock)
+    .replace(/\s+/g, '\\s*')
+    .replace(/>\s*</g, '>\\s*<')
+    .replace(/\s*>/g, '\\s*>');
+  
+  return new RegExp(searchRegex, 'g');
+};
+
+/**
+ * Process AI response chunk and apply updates to pages
+ * Returns updated pages and updated line numbers
+ */
+export interface ProcessAiResponseResult {
+  updatedPages: Page[];
+  updatedLines: number[][];
+}
+
+export const processAiResponse = (
+  chunk: string,
+  pages: Page[]
+): ProcessAiResponseResult => {
+  const updatedLines: number[][] = [];
+  const updatedPages = [...pages];
+
+  // Process UPDATE_FILE blocks
+  const updateFileRegex = new RegExp(
+    `${UPDATE_FILE_START.replace(/[.*+?^${}()|[\]\\]/g, '\\$&')}([^\\s]+)\\s*${UPDATE_FILE_END.replace(/[.*+?^${}()|[\]\\]/g, '\\$&')}([\\s\\S]*?)(?=${UPDATE_FILE_START.replace(/[.*+?^${}()|[\]\\]/g, '\\$&')}|${NEW_FILE_START.replace(/[.*+?^${}()|[\]\\]/g, '\\$&')}|$)`,
+    'g'
+  );
+  let updateFileMatch;
+
+  while ((updateFileMatch = updateFileRegex.exec(chunk)) !== null) {
+    const [, filePath, fileContent] = updateFileMatch;
+
+    const pageIndex = updatedPages.findIndex(p => p.path === filePath);
+    if (pageIndex !== -1) {
+      let pageHtml = updatedPages[pageIndex].html;
+
+      let processedContent = fileContent;
+      const htmlMatch = fileContent.match(/```html\s*([\s\S]*?)\s*```/);
+      if (htmlMatch) {
+        processedContent = htmlMatch[1];
+      }
+      let position = 0;
+      let moreBlocks = true;
+
+      while (moreBlocks) {
+        const searchStartIndex = processedContent.indexOf(SEARCH_START, position);
+        if (searchStartIndex === -1) {
+          moreBlocks = false;
+          continue;
+        }
+
+        const dividerIndex = processedContent.indexOf(DIVIDER, searchStartIndex);
+        if (dividerIndex === -1) {
+          moreBlocks = false;
+          continue;
+        }
+
+        const replaceEndIndex = processedContent.indexOf(REPLACE_END, dividerIndex);
+        if (replaceEndIndex === -1) {
+          moreBlocks = false;
+          continue;
+        }
+
+        const searchBlock = processedContent.substring(
+          searchStartIndex + SEARCH_START.length,
+          dividerIndex
+        );
+        const replaceBlock = processedContent.substring(
+          dividerIndex + DIVIDER.length,
+          replaceEndIndex
+        );
+
+        if (searchBlock.trim() === "") {
+          pageHtml = `${replaceBlock}\n${pageHtml}`;
+          updatedLines.push([1, replaceBlock.split("\n").length]);
+        } else {
+          const regex = createFlexibleHtmlRegex(searchBlock);
+          const match = regex.exec(pageHtml);
+
+          if (match) {
+            const matchedText = match[0];
+            const beforeText = pageHtml.substring(0, match.index);
+            const startLineNumber = beforeText.split("\n").length;
+            const replaceLines = replaceBlock.split("\n").length;
+            const endLineNumber = startLineNumber + replaceLines - 1;
+
+            updatedLines.push([startLineNumber, endLineNumber]);
+            pageHtml = pageHtml.replace(matchedText, replaceBlock);
+          }
+        }
+
+        position = replaceEndIndex + REPLACE_END.length;
+      }
+
+      updatedPages[pageIndex].html = pageHtml;
+    }
+  }
+
+  // Process NEW_FILE blocks
+  const newFileRegex = new RegExp(
+    `${NEW_FILE_START.replace(/[.*+?^${}()|[\]\\]/g, '\\$&')}([^\\s]+)\\s*${NEW_FILE_END.replace(/[.*+?^${}()|[\]\\]/g, '\\$&')}([\\s\\S]*?)(?=${UPDATE_FILE_START.replace(/[.*+?^${}()|[\]\\]/g, '\\$&')}|${NEW_FILE_START.replace(/[.*+?^${}()|[\]\\]/g, '\\$&')}|$)`,
+    'g'
+  );
+  let newFileMatch;
+
+  while ((newFileMatch = newFileRegex.exec(chunk)) !== null) {
+    const [, filePath, fileContent] = newFileMatch;
+
+    let fileData = fileContent;
+    // Try to extract content from code blocks
+    const htmlMatch = fileContent.match(/```html\s*([\s\S]*?)\s*```/);
+    const cssMatch = fileContent.match(/```css\s*([\s\S]*?)\s*```/);
+    const jsMatch = fileContent.match(/```javascript\s*([\s\S]*?)\s*```/);
+
+    if (htmlMatch) {
+      fileData = htmlMatch[1];
+    } else if (cssMatch) {
+      fileData = cssMatch[1];
+    } else if (jsMatch) {
+      fileData = jsMatch[1];
+    }
+
+    const existingFileIndex = updatedPages.findIndex(p => p.path === filePath);
+
+    if (existingFileIndex !== -1) {
+      updatedPages[existingFileIndex] = {
+        path: filePath,
+        html: fileData.trim()
+      };
+    } else {
+      updatedPages.push({
+        path: filePath,
+        html: fileData.trim()
+      });
+    }
+  }
+
+  // Fallback: process SEARCH/REPLACE blocks without UPDATE_FILE wrapper (backward compatibility)
+  if (updatedPages.length === pages.length && !chunk.includes(UPDATE_FILE_START)) {
+    let position = 0;
+    let moreBlocks = true;
+    let newHtml = updatedPages[0]?.html || "";
+
+    while (moreBlocks) {
+      const searchStartIndex = chunk.indexOf(SEARCH_START, position);
+      if (searchStartIndex === -1) {
+        moreBlocks = false;
+        continue;
+      }
+
+      const dividerIndex = chunk.indexOf(DIVIDER, searchStartIndex);
+      if (dividerIndex === -1) {
+        moreBlocks = false;
+        continue;
+      }
+
+      const replaceEndIndex = chunk.indexOf(REPLACE_END, dividerIndex);
+      if (replaceEndIndex === -1) {
+        moreBlocks = false;
+        continue;
+      }
+
+      const searchBlock = chunk.substring(
+        searchStartIndex + SEARCH_START.length,
+        dividerIndex
+      );
+      const replaceBlock = chunk.substring(
+        dividerIndex + DIVIDER.length,
+        replaceEndIndex
+      );
+
+      if (searchBlock.trim() === "") {
+        newHtml = `${replaceBlock}\n${newHtml}`;
+        updatedLines.push([1, replaceBlock.split("\n").length]);
+      } else {
+        const regex = createFlexibleHtmlRegex(searchBlock);
+        const match = regex.exec(newHtml);
+
+        if (match) {
+          const matchedText = match[0];
+          const beforeText = newHtml.substring(0, match.index);
+          const startLineNumber = beforeText.split("\n").length;
+          const replaceLines = replaceBlock.split("\n").length;
+          const endLineNumber = startLineNumber + replaceLines - 1;
+
+          updatedLines.push([startLineNumber, endLineNumber]);
+          newHtml = newHtml.replace(matchedText, replaceBlock);
+        }
+      }
+
+      position = replaceEndIndex + REPLACE_END.length;
+    }
+
+    const mainPageIndex = updatedPages.findIndex(p => p.path === '/' || p.path === '/index' || p.path === 'index');
+    if (mainPageIndex !== -1) {
+      updatedPages[mainPageIndex].html = newHtml;
+    }
+  }
+
+  return { updatedPages, updatedLines };
+};
+
+/**
+ * Convert pages to File objects for upload to HuggingFace
+ */
+export const pagesToFiles = (pages: Page[]): File[] => {
+  const files: File[] = [];
+  pages.forEach((page: Page) => {
+    let mimeType = "text/html";
+    if (page.path.endsWith(".css")) {
+      mimeType = "text/css";
+    } else if (page.path.endsWith(".js")) {
+      mimeType = "text/javascript";
+    } else if (page.path.endsWith(".json")) {
+      mimeType = "application/json";
+    }
+    const file = new File([page.html], page.path, { type: mimeType });
+    files.push(file);
+  });
+  return files;
+};
+
+/**
+ * Extract project name from AI response
+ */
+export const extractProjectName = (chunk: string): string | null => {
+  const projectName = chunk.match(/<<<<<<< PROJECT_NAME_START\s*([\s\S]*?)\s*>>>>>>> PROJECT_NAME_END/)?.[1]?.trim();
+  return projectName || null;
+};
+