refactor: move initialization code to SandboxManager

This commit is contained in:
James Murdza 2024-10-25 06:38:25 -06:00
parent 16e0c250d6
commit 3e891e6ab1
2 changed files with 107 additions and 113 deletions

View File

@ -1,8 +1,9 @@
import { Sandbox } from "e2b"
import { Socket } from 'socket.io' import { Socket } from 'socket.io'
import { AIWorker } from "./AIWorker" import { AIWorker } from "./AIWorker"
import { CONTAINER_TIMEOUT } from "./constants" import { CONTAINER_TIMEOUT } from "./constants"
import { DokkuClient } from "./DokkuClient" import { DokkuClient } from "./DokkuClient"
import { FileManager } from "./FileManager" import { FileManager, SandboxFiles } from "./FileManager"
import { import {
createFileRL, createFileRL,
createFolderRL, createFolderRL,
@ -14,6 +15,8 @@ import { SecureGitClient } from "./SecureGitClient"
import { TerminalManager } from "./TerminalManager" import { TerminalManager } from "./TerminalManager"
import { LockManager } from "./utils" import { LockManager } from "./utils"
const lockManager = new LockManager()
// Define a type for SocketHandler functions // Define a type for SocketHandler functions
type SocketHandler<T = Record<string, any>> = (args: T) => any; type SocketHandler<T = Record<string, any>> = (args: T) => any;
@ -25,53 +28,99 @@ function extractPortNumber(inputString: string): number | null {
return match ? parseInt(match[1]) : null return match ? parseInt(match[1]) : null
} }
export class SandboxManager { type SandboxManagerContext = {
fileManager: FileManager;
terminalManager: TerminalManager;
container: any;
aiWorker: AIWorker; aiWorker: AIWorker;
dokkuClient: DokkuClient | null; dokkuClient: DokkuClient | null;
gitClient: SecureGitClient | null; gitClient: SecureGitClient | null;
lockManager: LockManager;
socket: Socket; socket: Socket;
};
constructor(fileManager: FileManager, terminalManager: TerminalManager, aiWorker: AIWorker, dokkuClient: DokkuClient | null, gitClient: SecureGitClient | null, lockManager: LockManager, sandboxManager: any, socket: Socket) { export class SandboxManager {
this.fileManager = fileManager; fileManager: FileManager | null;
this.terminalManager = terminalManager; 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.aiWorker = aiWorker;
this.dokkuClient = dokkuClient; this.dokkuClient = dokkuClient;
this.gitClient = gitClient; this.gitClient = gitClient;
this.lockManager = lockManager;
this.socket = socket; 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.files)
}
}
async disconnect() {
await this.terminalManager?.closeAllTerminals()
await this.fileManager?.closeWatchers()
} }
handlers() { handlers() {
// Handle heartbeat from a socket connection // Handle heartbeat from a socket connection
const handleHeartbeat: SocketHandler = (_: any) => { const handleHeartbeat: SocketHandler = (_: any) => {
this.container.setTimeout(CONTAINER_TIMEOUT) this.container?.setTimeout(CONTAINER_TIMEOUT)
} }
// Handle getting a file // Handle getting a file
const handleGetFile: SocketHandler = ({ fileId }: any) => { const handleGetFile: SocketHandler = ({ fileId }: any) => {
return this.fileManager.getFile(fileId) return this.fileManager?.getFile(fileId)
} }
// Handle getting a folder // Handle getting a folder
const handleGetFolder: SocketHandler = ({ folderId }: any) => { const handleGetFolder: SocketHandler = ({ folderId }: any) => {
return this.fileManager.getFolder(folderId) return this.fileManager?.getFolder(folderId)
} }
// Handle saving a file // Handle saving a file
const handleSaveFile: SocketHandler = async ({ fileId, body, userId }: any) => { const handleSaveFile: SocketHandler = async ({ fileId, body }: any) => {
await saveFileRL.consume(userId, 1); await saveFileRL.consume(this.userId, 1);
return this.fileManager.saveFile(fileId, body) return this.fileManager?.saveFile(fileId, body)
} }
// Handle moving a file // Handle moving a file
const handleMoveFile: SocketHandler = ({ fileId, folderId }: any) => { const handleMoveFile: SocketHandler = ({ fileId, folderId }: any) => {
return this.fileManager.moveFile(fileId, folderId) return this.fileManager?.moveFile(fileId, folderId)
} }
// Handle listing apps // Handle listing apps
@ -81,55 +130,56 @@ export class SandboxManager {
} }
// Handle deploying code // Handle deploying code
const handleDeploy: SocketHandler = async ({ sandboxId }: any) => { const handleDeploy: SocketHandler = async (_: any) => {
if (!this.gitClient) throw Error("Failed to retrieve apps list: No git client") if (!this.gitClient) throw Error("No git client")
const fixedFilePaths = this.fileManager.sandboxFiles.fileData.map((file) => ({ if (!this.fileManager) throw Error("No file manager")
const fixedFilePaths = this.fileManager?.sandboxFiles.fileData.map((file) => ({
...file, ...file,
id: file.id.split("/").slice(2).join("/"), id: file.id.split("/").slice(2).join("/"),
})) }))
await this.gitClient.pushFiles(fixedFilePaths, sandboxId) await this.gitClient.pushFiles(fixedFilePaths, this.sandboxId)
return { success: true } return { success: true }
} }
// Handle creating a file // Handle creating a file
const handleCreateFile: SocketHandler = async ({ name, userId }: any) => { const handleCreateFile: SocketHandler = async ({ name }: any) => {
await createFileRL.consume(userId, 1); await createFileRL.consume(this.userId, 1);
return { "success": await this.fileManager.createFile(name) } return { "success": await this.fileManager?.createFile(name) }
} }
// Handle creating a folder // Handle creating a folder
const handleCreateFolder: SocketHandler = async ({ name, userId }: any) => { const handleCreateFolder: SocketHandler = async ({ name }: any) => {
await createFolderRL.consume(userId, 1); await createFolderRL.consume(this.userId, 1);
return { "success": await this.fileManager.createFolder(name) } return { "success": await this.fileManager?.createFolder(name) }
} }
// Handle renaming a file // Handle renaming a file
const handleRenameFile: SocketHandler = async ({ fileId, newName, userId }: any) => { const handleRenameFile: SocketHandler = async ({ fileId, newName }: any) => {
await renameFileRL.consume(userId, 1) await renameFileRL.consume(this.userId, 1)
return this.fileManager.renameFile(fileId, newName) return this.fileManager?.renameFile(fileId, newName)
} }
// Handle deleting a file // Handle deleting a file
const handleDeleteFile: SocketHandler = async ({ fileId, userId }: any) => { const handleDeleteFile: SocketHandler = async ({ fileId }: any) => {
await deleteFileRL.consume(userId, 1) await deleteFileRL.consume(this.userId, 1)
return this.fileManager.deleteFile(fileId) return this.fileManager?.deleteFile(fileId)
} }
// Handle deleting a folder // Handle deleting a folder
const handleDeleteFolder: SocketHandler = ({ folderId }: any) => { const handleDeleteFolder: SocketHandler = ({ folderId }: any) => {
return this.fileManager.deleteFolder(folderId) return this.fileManager?.deleteFolder(folderId)
} }
// Handle creating a terminal session // Handle creating a terminal session
const handleCreateTerminal: SocketHandler = async ({ id, sandboxId }: any) => { const handleCreateTerminal: SocketHandler = async ({ id }: any) => {
await this.lockManager.acquireLock(sandboxId, async () => { await lockManager.acquireLock(this.sandboxId, async () => {
await this.terminalManager.createTerminal(id, (responseString: string) => { await this.terminalManager?.createTerminal(id, (responseString: string) => {
this.socket.emit("terminalResponse", { id, data: responseString }) this.socket.emit("terminalResponse", { id, data: responseString })
const port = extractPortNumber(responseString) const port = extractPortNumber(responseString)
if (port) { if (port) {
this.socket.emit( this.socket.emit(
"previewURL", "previewURL",
"https://" + this.container.getHost(port) "https://" + this.container?.getHost(port)
) )
} }
}) })
@ -138,22 +188,22 @@ export class SandboxManager {
// Handle resizing a terminal // Handle resizing a terminal
const handleResizeTerminal: SocketHandler = ({ dimensions }: any) => { const handleResizeTerminal: SocketHandler = ({ dimensions }: any) => {
this.terminalManager.resizeTerminal(dimensions) this.terminalManager?.resizeTerminal(dimensions)
} }
// Handle sending data to a terminal // Handle sending data to a terminal
const handleTerminalData: SocketHandler = ({ id, data }: any) => { const handleTerminalData: SocketHandler = ({ id, data }: any) => {
return this.terminalManager.sendTerminalData(id, data) return this.terminalManager?.sendTerminalData(id, data)
} }
// Handle closing a terminal // Handle closing a terminal
const handleCloseTerminal: SocketHandler = ({ id }: any) => { const handleCloseTerminal: SocketHandler = ({ id }: any) => {
return this.terminalManager.closeTerminal(id) return this.terminalManager?.closeTerminal(id)
} }
// Handle generating code // Handle generating code
const handleGenerateCode: SocketHandler = ({ userId, fileName, code, line, instructions }: any) => { const handleGenerateCode: SocketHandler = ({ fileName, code, line, instructions }: any) => {
return this.aiWorker.generateCode(userId, fileName, code, line, instructions) return this.aiWorker.generateCode(this.userId, fileName, code, line, instructions)
} }
return { return {

View File

@ -1,19 +1,15 @@
import cors from "cors" import cors from "cors"
import dotenv from "dotenv" import dotenv from "dotenv"
import { Sandbox } from "e2b"
import express, { Express } from "express" import express, { Express } from "express"
import fs from "fs" import fs from "fs"
import { createServer } from "http" import { createServer } from "http"
import { Server } from "socket.io" import { Server } from "socket.io"
import { AIWorker } from "./AIWorker" import { AIWorker } from "./AIWorker"
import { CONTAINER_TIMEOUT } from "./constants"
import { DokkuClient } from "./DokkuClient" import { DokkuClient } from "./DokkuClient"
import { FileManager, SandboxFiles } from "./FileManager"
import { SandboxManager } from "./SandboxManager" import { SandboxManager } from "./SandboxManager"
import { SecureGitClient } from "./SecureGitClient" import { SecureGitClient } from "./SecureGitClient"
import { socketAuth } from "./socketAuth"; // Import the new socketAuth middleware import { socketAuth } from "./socketAuth"; // Import the new socketAuth middleware
import { TerminalManager } from "./TerminalManager"
import { LockManager } from "./utils"
// Handle uncaught exceptions // Handle uncaught exceptions
process.on("uncaughtException", (error) => { process.on("uncaughtException", (error) => {
@ -35,10 +31,8 @@ function isOwnerConnected(sandboxId: string): boolean {
} }
// Initialize containers and managers // Initialize containers and managers
const containers: Record<string, Sandbox> = {}
const connections: Record<string, number> = {} const connections: Record<string, number> = {}
const fileManagers: Record<string, FileManager> = {} const sandboxManagers: Record<string, SandboxManager> = {}
const terminalManagers: Record<string, TerminalManager> = {}
// Load environment variables // Load environment variables
dotenv.config() dotenv.config()
@ -57,9 +51,6 @@ const io = new Server(httpServer, {
// Middleware for socket authentication // Middleware for socket authentication
io.use(socketAuth) // Use the new socketAuth middleware io.use(socketAuth) // Use the new socketAuth middleware
// Initialize lock manager
const lockManager = new LockManager()
// Check for required environment variables // Check for required environment variables
if (!process.env.DOKKU_HOST) if (!process.env.DOKKU_HOST)
console.warn("Environment variable DOKKU_HOST is not defined") console.warn("Environment variable DOKKU_HOST is not defined")
@ -99,6 +90,7 @@ const aiWorker = new AIWorker(
// Handle a client connecting to the server // Handle a client connecting to the server
io.on("connection", async (socket) => { io.on("connection", async (socket) => {
try { try {
// This data comes is added by our authentication middleware
const data = socket.data as { const data = socket.data as {
userId: string userId: string
sandboxId: string sandboxId: string
@ -115,70 +107,23 @@ io.on("connection", async (socket) => {
} }
} }
// Create or retrieve container const sandboxManager = sandboxManagers[data.sandboxId] ?? new SandboxManager(
const createdContainer = await lockManager.acquireLock(
data.sandboxId, data.sandboxId,
async () => { data.userId,
{ aiWorker, dokkuClient, gitClient, socket }
)
try { try {
// Start a new container if the container doesn't exist or it timed out. sandboxManager.initializeContainer()
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) { } catch (e: any) {
console.error(`Error creating container ${data.sandboxId}:`, e) console.error(`Error initializing sandbox ${data.sandboxId}:`, e);
socket.emit("error", `Error: container creation. ${e.message ?? e}`) socket.emit("error", `Error: initialize sandbox ${data.sandboxId}. ${e.message ?? e}`);
} }
}
)
// Function to send loaded event
const sendLoadedEvent = (files: SandboxFiles) => {
socket.emit("loaded", files.files)
}
// 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]) => { Object.entries(sandboxManager.handlers()).forEach(([event, handler]) => {
socket.on(event, async (options: any, callback?: (response: any) => void) => { socket.on(event, async (options: any, callback?: (response: any) => void) => {
try { try {
// Consume rate limiter if provided const response = await handler(options)
const response = await handler({ ...options, ...data })
callback?.(response); callback?.(response);
} catch (e: any) { } catch (e: any) {
console.error(`Error processing event "${event}":`, e); console.error(`Error processing event "${event}":`, e);
@ -193,8 +138,7 @@ io.on("connection", async (socket) => {
connections[data.sandboxId]-- connections[data.sandboxId]--
} }
await terminalManager.closeAllTerminals() await sandboxManager.disconnect()
await fileManager.closeWatchers()
if (data.isOwner && connections[data.sandboxId] <= 0) { if (data.isOwner && connections[data.sandboxId] <= 0) {
socket.broadcast.emit( socket.broadcast.emit(