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 { 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<T = Record<string, any>> = (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.files)
}
}
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 {

View File

@ -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<string, Sandbox> = {}
const connections: Record<string, number> = {}
const fileManagers: Record<string, FileManager> = {}
const terminalManagers: Record<string, TerminalManager> = {}
const sandboxManagers: Record<string, SandboxManager> = {}
// 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(