improve ai logic

backend/server/src/index.ts

@@ -1,26 +1,25 @@
-import fs from "fs"
+import fs from "fs";
-import os from "os"
+import os from "os";
-import path from "path"
+import path from "path";
-import cors from "cors"
+import cors from "cors";
-import express, { Express } from "express"
+import express, { Express } from "express";
-import dotenv from "dotenv"
+import dotenv from "dotenv";
-import { createServer } from "http"
+import { createServer } from "http";
-import { Server } from "socket.io"
+import { Server } from "socket.io";
 
-import { z } from "zod"
+import { z } from "zod";
-import { User } from "./types"
+import { User } from "./types";
 import {
   createFile,
   deleteFile,
-  generateCode,
   getFolder,
   getProjectSize,
   getSandboxFiles,
   renameFile,
   saveFile,
   stopServer,
-} from "./utils"
+} from "./utils";
-import { IDisposable, IPty, spawn } from "node-pty"
+import { IDisposable, IPty, spawn } from "node-pty";
 import {
   MAX_BODY_SIZE,
   createFileRL,
@@ -28,356 +27,359 @@ import {
   deleteFileRL,
   renameFileRL,
   saveFileRL,
-} from "./ratelimit"
+} from "./ratelimit";
 
-dotenv.config()
+dotenv.config();
 
-const app: Express = express()
+const app: Express = express();
-const port = process.env.PORT || 4000
+const port = process.env.PORT || 4000;
-app.use(cors())
+app.use(cors());
-const httpServer = createServer(app)
+const httpServer = createServer(app);
 const io = new Server(httpServer, {
   cors: {
     origin: "*",
   },
-})
+});
 
 let inactivityTimeout: NodeJS.Timeout | null = null;
 let isOwnerConnected = false;
 
 const terminals: {
-  [id: string]: { terminal: IPty; onData: IDisposable; onExit: IDisposable }
+  [id: string]: { terminal: IPty; onData: IDisposable; onExit: IDisposable };
-} = {}
+} = {};
 
-const dirName = path.join(__dirname, "..")
+const dirName = path.join(__dirname, "..");
 
 const handshakeSchema = z.object({
   userId: z.string(),
   sandboxId: z.string(),
   EIO: z.string(),
   transport: z.string(),
-})
+});
 
 io.use(async (socket, next) => {
-  const q = socket.handshake.query
+  const q = socket.handshake.query;
-  const parseQuery = handshakeSchema.safeParse(q)
+  const parseQuery = handshakeSchema.safeParse(q);
 
   if (!parseQuery.success) {
-    ("Invalid request.")
+    ("Invalid request.");
-    next(new Error("Invalid request."))
+    next(new Error("Invalid request."));
-    return
+    return;
   }
 
-  const { sandboxId, userId } = parseQuery.data
+  const { sandboxId, userId } = parseQuery.data;
-  const dbUser = await fetch(`https://database.ishaan1013.workers.dev/api/user?id=${userId}`)
+  const dbUser = await fetch(
-  const dbUserJSON = (await dbUser.json()) as User
+    `https://database.ishaan1013.workers.dev/api/user?id=${userId}`
+  );
+  const dbUserJSON = (await dbUser.json()) as User;
 
   if (!dbUserJSON) {
-    next(new Error("DB error."))
+    next(new Error("DB error."));
-    return
+    return;
   }
 
-  const sandbox = dbUserJSON.sandbox.find((s) => s.id === sandboxId)
+  const sandbox = dbUserJSON.sandbox.find((s) => s.id === sandboxId);
   const sharedSandboxes = dbUserJSON.usersToSandboxes.find(
     (uts) => uts.sandboxId === sandboxId
-  )
+  );
 
   if (!sandbox && !sharedSandboxes) {
-    next(new Error("Invalid credentials."))
+    next(new Error("Invalid credentials."));
-    return
+    return;
   }
 
   socket.data = {
     userId,
     sandboxId: sandboxId,
     isOwner: sandbox !== undefined,
-  }
+  };
 
-  next()
+  next();
-})
+});
 
 io.on("connection", async (socket) => {
-
   if (inactivityTimeout) clearTimeout(inactivityTimeout);
 
   const data = socket.data as {
-    userId: string
+    userId: string;
-    sandboxId: string
+    sandboxId: string;
-    isOwner: boolean
+    isOwner: boolean;
-  }
+  };
 
   if (data.isOwner) {
-    isOwnerConnected = true
+    isOwnerConnected = true;
   } else {
     if (!isOwnerConnected) {
-      socket.emit("disableAccess", "The sandbox owner is not connected.")
+      socket.emit("disableAccess", "The sandbox owner is not connected.");
-      return
+      return;
     }
   }
 
-  const sandboxFiles = await getSandboxFiles(data.sandboxId)
+  const sandboxFiles = await getSandboxFiles(data.sandboxId);
   sandboxFiles.fileData.forEach((file) => {
-    const filePath = path.join(dirName, file.id)
+    const filePath = path.join(dirName, file.id);
-    fs.mkdirSync(path.dirname(filePath), { recursive: true })
+    fs.mkdirSync(path.dirname(filePath), { recursive: true });
     fs.writeFile(filePath, file.data, function (err) {
-      if (err) throw err
+      if (err) throw err;
-    })
+    });
-  })
+  });
 
-  socket.emit("loaded", sandboxFiles.files)
+  socket.emit("loaded", sandboxFiles.files);
 
   socket.on("getFile", (fileId: string, callback) => {
-    const file = sandboxFiles.fileData.find((f) => f.id === fileId)
+    const file = sandboxFiles.fileData.find((f) => f.id === fileId);
-    if (!file) return
+    if (!file) return;
 
-    callback(file.data)
+    callback(file.data);
-  })
+  });
 
   socket.on("getFolder", async (folderId: string, callback) => {
-    const files = await getFolder(folderId)
+    const files = await getFolder(folderId);
-    callback(files)
+    callback(files);
-  })
+  });
 
   // todo: send diffs + debounce for efficiency
   socket.on("saveFile", async (fileId: string, body: string) => {
     try {
-      await saveFileRL.consume(data.userId, 1)
+      await saveFileRL.consume(data.userId, 1);
 
       if (Buffer.byteLength(body, "utf-8") > MAX_BODY_SIZE) {
         socket.emit(
           "rateLimit",
           "Rate limited: file size too large. Please reduce the file size."
-        )
+        );
-        return
+        return;
       }
 
-      const file = sandboxFiles.fileData.find((f) => f.id === fileId)
+      const file = sandboxFiles.fileData.find((f) => f.id === fileId);
-      if (!file) return
+      if (!file) return;
-      file.data = body
+      file.data = body;
 
       fs.writeFile(path.join(dirName, file.id), body, function (err) {
-        if (err) throw err
+        if (err) throw err;
-      })
+      });
-      await saveFile(fileId, body)
+      await saveFile(fileId, body);
     } catch (e) {
-      io.emit("rateLimit", "Rate limited: file saving. Please slow down.")
+      io.emit("rateLimit", "Rate limited: file saving. Please slow down.");
     }
-  })
+  });
 
   socket.on("moveFile", async (fileId: string, folderId: string, callback) => {
-    const file = sandboxFiles.fileData.find((f) => f.id === fileId)
+    const file = sandboxFiles.fileData.find((f) => f.id === fileId);
-    if (!file) return
+    if (!file) return;
 
-    const parts = fileId.split("/")
+    const parts = fileId.split("/");
-    const newFileId = folderId + "/" + parts.pop()
+    const newFileId = folderId + "/" + parts.pop();
 
     fs.rename(
       path.join(dirName, fileId),
       path.join(dirName, newFileId),
       function (err) {
-        if (err) throw err
+        if (err) throw err;
       }
-    )
+    );
 
-    file.id = newFileId
+    file.id = newFileId;
 
-    await renameFile(fileId, newFileId, file.data)
+    await renameFile(fileId, newFileId, file.data);
-    const newFiles = await getSandboxFiles(data.sandboxId)
+    const newFiles = await getSandboxFiles(data.sandboxId);
 
-    callback(newFiles.files)
+    callback(newFiles.files);
-  })
+  });
 
   socket.on("createFile", async (name: string, callback) => {
     try {
-
+      const size: number = await getProjectSize(data.sandboxId);
-      const size: number = await getProjectSize(data.sandboxId)
       // limit is 200mb
       if (size > 200 * 1024 * 1024) {
-        io.emit("rateLimit", "Rate limited: project size exceeded. Please delete some files.")
+        io.emit(
-        callback({success: false})
+          "rateLimit",
+          "Rate limited: project size exceeded. Please delete some files."
+        );
+        callback({ success: false });
       }
 
-      await createFileRL.consume(data.userId, 1)
+      await createFileRL.consume(data.userId, 1);
 
-      const id = `projects/${data.sandboxId}/${name}`
+      const id = `projects/${data.sandboxId}/${name}`;
 
       fs.writeFile(path.join(dirName, id), "", function (err) {
-        if (err) throw err
+        if (err) throw err;
-      })
+      });
 
       sandboxFiles.files.push({
         id,
         name,
         type: "file",
-      })
+      });
 
       sandboxFiles.fileData.push({
         id,
         data: "",
-      })
+      });
 
-      await createFile(id)
+      await createFile(id);
 
-      callback({success: true})
+      callback({ success: true });
     } catch (e) {
-      io.emit("rateLimit", "Rate limited: file creation. Please slow down.")
+      io.emit("rateLimit", "Rate limited: file creation. Please slow down.");
     }
-  })
+  });
 
   socket.on("createFolder", async (name: string, callback) => {
     try {
-      await createFolderRL.consume(data.userId, 1)
+      await createFolderRL.consume(data.userId, 1);
 
-      const id = `projects/${data.sandboxId}/${name}`
+      const id = `projects/${data.sandboxId}/${name}`;
 
       fs.mkdir(path.join(dirName, id), { recursive: true }, function (err) {
-        if (err) throw err
+        if (err) throw err;
-      })
+      });
 
-      callback()
+      callback();
     } catch (e) {
-      io.emit("rateLimit", "Rate limited: folder creation. Please slow down.")
+      io.emit("rateLimit", "Rate limited: folder creation. Please slow down.");
     }
-  })
+  });
 
   socket.on("renameFile", async (fileId: string, newName: string) => {
     try {
-      await renameFileRL.consume(data.userId, 1)
+      await renameFileRL.consume(data.userId, 1);
 
-      const file = sandboxFiles.fileData.find((f) => f.id === fileId)
+      const file = sandboxFiles.fileData.find((f) => f.id === fileId);
-      if (!file) return
+      if (!file) return;
-      file.id = newName
+      file.id = newName;
 
-      const parts = fileId.split("/")
+      const parts = fileId.split("/");
       const newFileId =
-        parts.slice(0, parts.length - 1).join("/") + "/" + newName
+        parts.slice(0, parts.length - 1).join("/") + "/" + newName;
 
       fs.rename(
         path.join(dirName, fileId),
         path.join(dirName, newFileId),
         function (err) {
-          if (err) throw err
+          if (err) throw err;
         }
-      )
+      );
-      await renameFile(fileId, newFileId, file.data)
+      await renameFile(fileId, newFileId, file.data);
     } catch (e) {
-      io.emit("rateLimit", "Rate limited: file renaming. Please slow down.")
+      io.emit("rateLimit", "Rate limited: file renaming. Please slow down.");
-      return
+      return;
     }
-  })
+  });
 
   socket.on("deleteFile", async (fileId: string, callback) => {
     try {
-      await deleteFileRL.consume(data.userId, 1)
+      await deleteFileRL.consume(data.userId, 1);
-      const file = sandboxFiles.fileData.find((f) => f.id === fileId)
+      const file = sandboxFiles.fileData.find((f) => f.id === fileId);
-      if (!file) return
+      if (!file) return;
 
       fs.unlink(path.join(dirName, fileId), function (err) {
-        if (err) throw err
+        if (err) throw err;
-      })
+      });
       sandboxFiles.fileData = sandboxFiles.fileData.filter(
         (f) => f.id !== fileId
-      )
+      );
 
-      await deleteFile(fileId)
+      await deleteFile(fileId);
 
-      const newFiles = await getSandboxFiles(data.sandboxId)
+      const newFiles = await getSandboxFiles(data.sandboxId);
-      callback(newFiles.files)
+      callback(newFiles.files);
     } catch (e) {
-      io.emit("rateLimit", "Rate limited: file deletion. Please slow down.")
+      io.emit("rateLimit", "Rate limited: file deletion. Please slow down.");
     }
-  })
+  });
 
   socket.on("renameFolder", async (folderId: string, newName: string) => {
     // todo
-  })
+  });
 
   socket.on("deleteFolder", async (folderId: string, callback) => {
-    const files = await getFolder(folderId)
+    const files = await getFolder(folderId);
 
-    await Promise.all(files.map(async (file) => {
+    await Promise.all(
+      files.map(async (file) => {
         fs.unlink(path.join(dirName, file), function (err) {
-        if (err) throw err
+          if (err) throw err;
-      })
+        });
 
         sandboxFiles.fileData = sandboxFiles.fileData.filter(
           (f) => f.id !== file
-      )
+        );
-
-      await deleteFile(file)
-    }))
-
-    const newFiles = await getSandboxFiles(data.sandboxId)
-    
-    callback(newFiles.files)
 
+        await deleteFile(file);
       })
+    );
+
+    const newFiles = await getSandboxFiles(data.sandboxId);
+
+    callback(newFiles.files);
+  });
 
   socket.on("createTerminal", (id: string, callback) => {
-    console.log("creating terminal", id)
+    console.log("creating terminal", id);
     if (terminals[id] || Object.keys(terminals).length >= 4) {
-      return
+      return;
     }
 
     const pty = spawn(os.platform() === "win32" ? "cmd.exe" : "bash", [], {
       name: "xterm",
       cols: 100,
       cwd: path.join(dirName, "projects", data.sandboxId),
-    })
+    });
 
     const onData = pty.onData((data) => {
       // console.log("terminalResponse", id, data)
       io.emit("terminalResponse", {
         id,
         data,
-      })
+      });
-    })
+    });
 
-    const onExit = pty.onExit((code) => console.log("exit :(", code))
+    const onExit = pty.onExit((code) => console.log("exit :(", code));
 
-    pty.write("clear\r")
+    pty.write("clear\r");
 
     terminals[id] = {
       terminal: pty,
       onData,
       onExit,
-    }
+    };
 
-    callback()
+    callback();
-  })
+  });
 
   socket.on("resizeTerminal", (dimensions: { cols: number; rows: number }) => {
-    console.log("resizeTerminal", dimensions)
+    console.log("resizeTerminal", dimensions);
     Object.values(terminals).forEach((t) => {
-      t.terminal.resize(dimensions.cols, dimensions.rows)
+      t.terminal.resize(dimensions.cols, dimensions.rows);
-    })
+    });
-  
+  });
-  })
 
   socket.on("terminalData", (id: string, data: string) => {
     if (!terminals[id]) {
-      console.log("terminal not found", id)
+      console.log("terminal not found", id);
-      return
+      return;
     }
 
     try {
-      terminals[id].terminal.write(data)
+      terminals[id].terminal.write(data);
     } catch (e) {
-      console.log("Error writing to terminal", e)
+      console.log("Error writing to terminal", e);
     }
-  })
+  });
 
   socket.on("closeTerminal", (id: string, callback) => {
     if (!terminals[id]) {
-      return
+      return;
     }
 
-    terminals[id].onData.dispose()
+    terminals[id].onData.dispose();
-    terminals[id].onExit.dispose()
+    terminals[id].onExit.dispose();
-    delete terminals[id]
+    delete terminals[id];
 
-    callback()
+    callback();
-  })
+  });
 
   socket.on(
     "generateCode",
@@ -389,7 +391,9 @@ io.on("connection", async (socket) => {
       callback
     ) => {
       // Log code generation credit in DB
-      const fetchPromise = fetch(`https://database.ishaan1013.workers.dev/api/sandbox/generate`, {
+      const fetchPromise = fetch(
+        `https://database.ishaan1013.workers.dev/api/sandbox/generate`,
+        {
           method: "POST",
           headers: {
             "Content-Type": "application/json",
@@ -397,62 +401,70 @@ io.on("connection", async (socket) => {
           body: JSON.stringify({
             userId: data.userId,
           }),
-      })
+        }
+      );
+
+      console.log("CF_AI_KEY", process.env.CF_AI_KEY);
 
       // Generate code from cloudflare workers AI
-      const generateCodePromise = generateCode({
+      const generateCodePromise = fetch(
-        fileName,
+        `https://ai.ishaan1013.workers.dev/api?fileName=${fileName}&code=${code}&line=${line}&instructions=${instructions}`,
-        code,
+        {
-        line,
+          headers: {
-        instructions,
+            "Content-Type": "application/json",
-      })
+            Authorization: `${process.env.CF_AI_KEY}`,
+          },
+        }
+      );
 
       const [fetchResponse, generateCodeResponse] = await Promise.all([
         fetchPromise,
         generateCodePromise,
-      ])
+      ]);
 
-      const json = await generateCodeResponse.json()
+      const json = await generateCodeResponse.json();
-      callback(json)
+      callback(json);
     }
-  )
+  );
 
   socket.on("disconnect", async () => {
-    console.log("disconnected", data.userId, data.sandboxId)
+    console.log("disconnected", data.userId, data.sandboxId);
     if (data.isOwner) {
       // console.log("deleting all terminals")
       Object.entries(terminals).forEach((t) => {
-        const { terminal, onData, onExit } = t[1]
+        const { terminal, onData, onExit } = t[1];
-        onData.dispose()
+        onData.dispose();
-        onExit.dispose()
+        onExit.dispose();
-        delete terminals[t[0]]
+        delete terminals[t[0]];
-      })
+      });
 
-      socket.broadcast.emit("disableAccess", "The sandbox owner has disconnected.")
+      socket.broadcast.emit(
+        "disableAccess",
+        "The sandbox owner has disconnected."
+      );
     }
 
-    const sockets = await io.fetchSockets()
+    const sockets = await io.fetchSockets();
     if (inactivityTimeout) {
-      clearTimeout(inactivityTimeout)
+      clearTimeout(inactivityTimeout);
-    };
+    }
     if (sockets.length === 0) {
-      console.log("STARTING TIMER")
+      console.log("STARTING TIMER");
       inactivityTimeout = setTimeout(() => {
         io.fetchSockets().then(async (sockets) => {
           if (sockets.length === 0) {
             // close server
             console.log("Closing server due to inactivity.");
-            const res = await stopServer(data.sandboxId, data.userId)
+            const res = await stopServer(data.sandboxId, data.userId);
           }
         });
       }, 20000);
     } else {
-      console.log("number of sockets", sockets.length)
+      console.log("number of sockets", sockets.length);
     }
-
+  });
-  })
+});
-})
 
 httpServer.listen(port, () => {
-  console.log(`Server, running on port ${port}`)
+  console.log(`Server, running on port ${port}`);
-})
+});

backend/server/src/utils.ts

@@ -151,28 +151,6 @@ export const getProjectSize = async (id: string) => {
   return (await res.json()).size;
 };
 
-export const generateCode = async ({
-  fileName,
-  code,
-  line,
-  instructions,
-}: {
-  fileName: string;
-  code: string;
-  line: number;
-  instructions: string;
-}) => {
-  return await fetch(
-    `https://ai.ishaan1013.workers.dev/api?fileName=${fileName}&code=${code}&line=${line}&instructions=${instructions}`,
-    {
-      headers: {
-        "Content-Type": "application/json",
-        Authorization: `${process.env.CF_AI_KEY}`,
-      },
-    }
-  );
-};
-
 export const stopServer = async (sandboxId: string, userId: string) => {
   const res = await fetch("http://localhost:4001/stop", {
     method: "POST",

frontend/components/editor/generate.tsx

@@ -49,8 +49,8 @@ export default function GenerateInput({
   useEffect(() => {
     setTimeout(() => {
       inputRef.current?.focus();
-    }, 0);
+    }, 100);
-  }, []);
+  }, [inputRef.current]);
 
   const handleGenerate = async ({
     regenerate = false,