refactor: move initialization code to SandboxManager
This commit is contained in:
parent
15fbd4ce41
commit
a6f457ef59
@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
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 {
|
||||
|
@ -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 () => {
|
||||
data.userId,
|
||||
{ aiWorker, dokkuClient, gitClient, socket }
|
||||
)
|
||||
|
||||
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
|
||||
}
|
||||
sandboxManager.initializeContainer()
|
||||
} catch (e: any) {
|
||||
console.error(`Error creating container ${data.sandboxId}:`, e)
|
||||
socket.emit("error", `Error: container creation. ${e.message ?? e}`)
|
||||
console.error(`Error initializing sandbox ${data.sandboxId}:`, 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]) => {
|
||||
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(
|
||||
|
Loading…
x
Reference in New Issue
Block a user