From a6f457ef59889d471c68cc0a0a035ac7ae063066 Mon Sep 17 00:00:00 2001 From: James Murdza Date: Fri, 25 Oct 2024 06:38:25 -0600 Subject: [PATCH] refactor: move initialization code to SandboxManager --- backend/server/src/SandboxManager.ts | 138 ++++++++++++++++++--------- backend/server/src/index.ts | 82 +++------------- 2 files changed, 107 insertions(+), 113 deletions(-) diff --git a/backend/server/src/SandboxManager.ts b/backend/server/src/SandboxManager.ts index 4d957bf..d64a816 100644 --- a/backend/server/src/SandboxManager.ts +++ b/backend/server/src/SandboxManager.ts @@ -1,8 +1,9 @@ +import { Sandbox } from "e2b" import { Socket } from 'socket.io' import { AIWorker } from "./AIWorker" import { CONTAINER_TIMEOUT } from "./constants" import { DokkuClient } from "./DokkuClient" -import { FileManager } from "./FileManager" +import { FileManager, SandboxFiles } from "./FileManager" import { createFileRL, createFolderRL, @@ -14,6 +15,8 @@ import { SecureGitClient } from "./SecureGitClient" import { TerminalManager } from "./TerminalManager" import { LockManager } from "./utils" +const lockManager = new LockManager() + // Define a type for SocketHandler functions type SocketHandler> = (args: T) => any; @@ -25,53 +28,99 @@ function extractPortNumber(inputString: string): number | null { return match ? parseInt(match[1]) : null } -export class SandboxManager { - fileManager: FileManager; - terminalManager: TerminalManager; - container: any; +type SandboxManagerContext = { aiWorker: AIWorker; dokkuClient: DokkuClient | null; gitClient: SecureGitClient | null; - lockManager: LockManager; socket: Socket; +}; - constructor(fileManager: FileManager, terminalManager: TerminalManager, aiWorker: AIWorker, dokkuClient: DokkuClient | null, gitClient: SecureGitClient | null, lockManager: LockManager, sandboxManager: any, socket: Socket) { - this.fileManager = fileManager; - this.terminalManager = terminalManager; +export class SandboxManager { + fileManager: FileManager | null; + terminalManager: TerminalManager | null; + container: Sandbox | null; + dokkuClient: DokkuClient | null; + gitClient: SecureGitClient | null; + aiWorker: AIWorker; + socket: Socket; + sandboxId: string; + userId: string; + + constructor(sandboxId: string, userId: string, { aiWorker, dokkuClient, gitClient, socket }: SandboxManagerContext) { + this.fileManager = null; + this.terminalManager = null; + this.container = null; + this.sandboxId = sandboxId; + this.userId = userId; this.aiWorker = aiWorker; this.dokkuClient = dokkuClient; this.gitClient = gitClient; - this.lockManager = lockManager; this.socket = socket; - this.container = sandboxManager; + } + + async initializeContainer() { + + await lockManager.acquireLock(this.sandboxId, async () => { + if (this.container && await this.container.isRunning()) { + console.log(`Found existing container ${this.sandboxId}`) + } else { + console.log("Creating container", this.sandboxId) + this.container = await Sandbox.create({ + timeoutMs: CONTAINER_TIMEOUT, + }) + } + }) + if (!this.container) throw new Error("Failed to create container") + + if (!this.terminalManager) { + this.terminalManager = new TerminalManager(this.container) + console.log(`Terminal manager set up for ${this.sandboxId}`) + } + + if (!this.fileManager) { + this.fileManager = new FileManager( + this.sandboxId, + this.container, + (files: SandboxFiles) => { + this.socket.emit("loaded", files.files) + } + ) + this.fileManager.initialize() + this.socket.emit("loaded", this.fileManager.sandboxFiles) + } + } + + async disconnect() { + await this.terminalManager?.closeAllTerminals() + await this.fileManager?.closeWatchers() } handlers() { // Handle heartbeat from a socket connection const handleHeartbeat: SocketHandler = (_: any) => { - this.container.setTimeout(CONTAINER_TIMEOUT) + this.container?.setTimeout(CONTAINER_TIMEOUT) } // Handle getting a file const handleGetFile: SocketHandler = ({ fileId }: any) => { - return this.fileManager.getFile(fileId) + return this.fileManager?.getFile(fileId) } // Handle getting a folder const handleGetFolder: SocketHandler = ({ folderId }: any) => { - return this.fileManager.getFolder(folderId) + return this.fileManager?.getFolder(folderId) } // Handle saving a file - const handleSaveFile: SocketHandler = async ({ fileId, body, userId }: any) => { - await saveFileRL.consume(userId, 1); - return this.fileManager.saveFile(fileId, body) + const handleSaveFile: SocketHandler = async ({ fileId, body }: any) => { + await saveFileRL.consume(this.userId, 1); + return this.fileManager?.saveFile(fileId, body) } // Handle moving a file const handleMoveFile: SocketHandler = ({ fileId, folderId }: any) => { - return this.fileManager.moveFile(fileId, folderId) + return this.fileManager?.moveFile(fileId, folderId) } // Handle listing apps @@ -81,55 +130,56 @@ export class SandboxManager { } // Handle deploying code - const handleDeploy: SocketHandler = async ({ sandboxId }: any) => { - if (!this.gitClient) throw Error("Failed to retrieve apps list: No git client") - const fixedFilePaths = this.fileManager.sandboxFiles.fileData.map((file) => ({ + const handleDeploy: SocketHandler = async (_: any) => { + if (!this.gitClient) throw Error("No git client") + if (!this.fileManager) throw Error("No file manager") + const fixedFilePaths = this.fileManager?.sandboxFiles.fileData.map((file) => ({ ...file, id: file.id.split("/").slice(2).join("/"), })) - await this.gitClient.pushFiles(fixedFilePaths, sandboxId) + await this.gitClient.pushFiles(fixedFilePaths, this.sandboxId) return { success: true } } // Handle creating a file - const handleCreateFile: SocketHandler = async ({ name, userId }: any) => { - await createFileRL.consume(userId, 1); - return { "success": await this.fileManager.createFile(name) } + const handleCreateFile: SocketHandler = async ({ name }: any) => { + await createFileRL.consume(this.userId, 1); + return { "success": await this.fileManager?.createFile(name) } } // Handle creating a folder - const handleCreateFolder: SocketHandler = async ({ name, userId }: any) => { - await createFolderRL.consume(userId, 1); - return { "success": await this.fileManager.createFolder(name) } + const handleCreateFolder: SocketHandler = async ({ name }: any) => { + await createFolderRL.consume(this.userId, 1); + return { "success": await this.fileManager?.createFolder(name) } } // Handle renaming a file - const handleRenameFile: SocketHandler = async ({ fileId, newName, userId }: any) => { - await renameFileRL.consume(userId, 1) - return this.fileManager.renameFile(fileId, newName) + const handleRenameFile: SocketHandler = async ({ fileId, newName }: any) => { + await renameFileRL.consume(this.userId, 1) + return this.fileManager?.renameFile(fileId, newName) } // Handle deleting a file - const handleDeleteFile: SocketHandler = async ({ fileId, userId }: any) => { - await deleteFileRL.consume(userId, 1) - return this.fileManager.deleteFile(fileId) + const handleDeleteFile: SocketHandler = async ({ fileId }: any) => { + await deleteFileRL.consume(this.userId, 1) + return this.fileManager?.deleteFile(fileId) } // Handle deleting a folder const handleDeleteFolder: SocketHandler = ({ folderId }: any) => { - return this.fileManager.deleteFolder(folderId) + return this.fileManager?.deleteFolder(folderId) } // Handle creating a terminal session - const handleCreateTerminal: SocketHandler = async ({ id, sandboxId }: any) => { - await this.lockManager.acquireLock(sandboxId, async () => { - await this.terminalManager.createTerminal(id, (responseString: string) => { + const handleCreateTerminal: SocketHandler = async ({ id }: any) => { + await lockManager.acquireLock(this.sandboxId, async () => { + await this.terminalManager?.createTerminal(id, (responseString: string) => { this.socket.emit("terminalResponse", { id, data: responseString }) const port = extractPortNumber(responseString) if (port) { this.socket.emit( "previewURL", - "https://" + this.container.getHost(port) + "https://" + this.container?.getHost(port) ) } }) @@ -138,22 +188,22 @@ export class SandboxManager { // Handle resizing a terminal const handleResizeTerminal: SocketHandler = ({ dimensions }: any) => { - this.terminalManager.resizeTerminal(dimensions) + this.terminalManager?.resizeTerminal(dimensions) } // Handle sending data to a terminal const handleTerminalData: SocketHandler = ({ id, data }: any) => { - return this.terminalManager.sendTerminalData(id, data) + return this.terminalManager?.sendTerminalData(id, data) } // Handle closing a terminal const handleCloseTerminal: SocketHandler = ({ id }: any) => { - return this.terminalManager.closeTerminal(id) + return this.terminalManager?.closeTerminal(id) } // Handle generating code - const handleGenerateCode: SocketHandler = ({ userId, fileName, code, line, instructions }: any) => { - return this.aiWorker.generateCode(userId, fileName, code, line, instructions) + const handleGenerateCode: SocketHandler = ({ fileName, code, line, instructions }: any) => { + return this.aiWorker.generateCode(this.userId, fileName, code, line, instructions) } return { diff --git a/backend/server/src/index.ts b/backend/server/src/index.ts index bdbc7af..54f1315 100644 --- a/backend/server/src/index.ts +++ b/backend/server/src/index.ts @@ -1,19 +1,15 @@ import cors from "cors" import dotenv from "dotenv" -import { Sandbox } from "e2b" import express, { Express } from "express" import fs from "fs" import { createServer } from "http" import { Server } from "socket.io" import { AIWorker } from "./AIWorker" -import { CONTAINER_TIMEOUT } from "./constants" + import { DokkuClient } from "./DokkuClient" -import { FileManager, SandboxFiles } from "./FileManager" import { SandboxManager } from "./SandboxManager" import { SecureGitClient } from "./SecureGitClient" import { socketAuth } from "./socketAuth"; // Import the new socketAuth middleware -import { TerminalManager } from "./TerminalManager" -import { LockManager } from "./utils" // Handle uncaught exceptions process.on("uncaughtException", (error) => { @@ -35,10 +31,8 @@ function isOwnerConnected(sandboxId: string): boolean { } // Initialize containers and managers -const containers: Record = {} const connections: Record = {} -const fileManagers: Record = {} -const terminalManagers: Record = {} +const sandboxManagers: Record = {} // Load environment variables dotenv.config() @@ -57,9 +51,6 @@ const io = new Server(httpServer, { // Middleware for socket authentication io.use(socketAuth) // Use the new socketAuth middleware -// Initialize lock manager -const lockManager = new LockManager() - // Check for required environment variables if (!process.env.DOKKU_HOST) console.warn("Environment variable DOKKU_HOST is not defined") @@ -99,6 +90,7 @@ const aiWorker = new AIWorker( // Handle a client connecting to the server io.on("connection", async (socket) => { try { + // This data comes is added by our authentication middleware const data = socket.data as { userId: string sandboxId: string @@ -115,70 +107,23 @@ io.on("connection", async (socket) => { } } - // Create or retrieve container - const createdContainer = await lockManager.acquireLock( + const sandboxManager = sandboxManagers[data.sandboxId] ?? new SandboxManager( data.sandboxId, - async () => { - try { - // Start a new container if the container doesn't exist or it timed out. - if ( - !containers[data.sandboxId] || - !(await containers[data.sandboxId].isRunning()) - ) { - containers[data.sandboxId] = await Sandbox.create({ - timeoutMs: CONTAINER_TIMEOUT, - }) - console.log("Created container ", data.sandboxId) - return true - } - } catch (e: any) { - console.error(`Error creating container ${data.sandboxId}:`, e) - socket.emit("error", `Error: container creation. ${e.message ?? e}`) - } - } + data.userId, + { aiWorker, dokkuClient, gitClient, socket } ) - // Function to send loaded event - const sendLoadedEvent = (files: SandboxFiles) => { - socket.emit("loaded", files.files) + try { + sandboxManager.initializeContainer() + } catch (e: any) { + console.error(`Error initializing sandbox ${data.sandboxId}:`, e); + socket.emit("error", `Error: initialize sandbox ${data.sandboxId}. ${e.message ?? e}`); } - // Initialize file and terminal managers if container was created - if (createdContainer) { - fileManagers[data.sandboxId] = new FileManager( - data.sandboxId, - containers[data.sandboxId], - sendLoadedEvent - ) - terminalManagers[data.sandboxId] = new TerminalManager( - containers[data.sandboxId] - ) - console.log(`terminal manager set up for ${data.sandboxId}`) - await fileManagers[data.sandboxId].initialize() - } - - const fileManager = fileManagers[data.sandboxId] - const terminalManager = terminalManagers[data.sandboxId] - - // Load file list from the file manager into the editor - sendLoadedEvent(fileManager.sandboxFiles) - - const sandboxManager = new SandboxManager( - fileManager, - terminalManager, - aiWorker, - dokkuClient, - gitClient, - lockManager, - containers[data.sandboxId], - socket - ) - Object.entries(sandboxManager.handlers()).forEach(([event, handler]) => { socket.on(event, async (options: any, callback?: (response: any) => void) => { try { - // Consume rate limiter if provided - const response = await handler({ ...options, ...data }) + const response = await handler(options) callback?.(response); } catch (e: any) { console.error(`Error processing event "${event}":`, e); @@ -193,8 +138,7 @@ io.on("connection", async (socket) => { connections[data.sandboxId]-- } - await terminalManager.closeAllTerminals() - await fileManager.closeWatchers() + await sandboxManager.disconnect() if (data.isOwner && connections[data.sandboxId] <= 0) { socket.broadcast.emit(