Spaces:

E5K7
/

asipilot

Sleeping

asipilot / packages /server /src /websocket /execution.ts

Update terminal, chat, and websocket functionality

4e105ef 7 months ago

7.55 kB

import { randomUUID } from 'crypto';
import type { Server as SocketServer } from 'socket.io';
import { exec, spawn, type ChildProcess } from 'child_process';
import { promises as fs } from 'fs';
import path from 'path';
import os from 'os';
import { WS_EVENTS } from '@asipilot/shared';
import { logger } from '../utils/logger.js';
import type { WSExecuteRequest } from '@asipilot/shared';
import { getStableSessionId } from './session-id.js';

/** Interactive runs may sit on input(); allow long thought + typing. */
const EXECUTION_TIMEOUT_MS = 120_000;

/** Child processes keyed by per-run id (client sends this with stdin). */
const activeChildProcs = new Map<string, ChildProcess>();
/** Latest run id for a browser tab session (cleared when run ends or new run starts). */
const sessionActiveRunId = new Map<string, string>();

const graceKillTimers = new Map<string, ReturnType<typeof setTimeout>>();

export function cancelGraceKill(sessionId: string): void {
  const t = graceKillTimers.get(sessionId);
  if (t !== undefined) {
    clearTimeout(t);
    graceKillTimers.delete(sessionId);
  }
}

export function scheduleGraceKill(sessionId: string, delayMs = 15000): void {
  cancelGraceKill(sessionId);
  const t = setTimeout(() => {
    graceKillTimers.delete(sessionId);
    cleanupProcess(sessionId);
    logger.debug(`Grace-killed runner for session ${sessionId} after disconnect`);
  }, delayMs);
  graceKillTimers.set(sessionId, t);
}

/** Deliver to every socket whose handshake matches this tab session (rooms alone can fail behind some proxies). */
function emitExec(io: SocketServer, sessionId: string, event: string, payload: object) {
  let n = 0;
  for (const socket of io.sockets.sockets.values()) {
    if (getStableSessionId(socket) === sessionId) {
      socket.emit(event, payload);
      n++;
    }
  }
  if (n === 0) {
    logger.warn(`exec ${event}: no socket matched session ${sessionId}; using room fallback`);
    io.to(sessionId).emit(event, payload);
  }
}

export async function handleCodeExecution(io: SocketServer, sessionId: string, data: WSExecuteRequest) {
  const { language, content } = data;
  let tempDir: string | null = null;

  cancelGraceKill(sessionId);
  cleanupProcess(sessionId);

  try {
    tempDir = await fs.mkdtemp(path.join(os.tmpdir(), 'asipilot-exec-'));

    if (language === 'python') {
      if (!tempDir) throw new Error('Temp dir not created');
      const dir = tempDir;
      const filePath = path.join(dir, 'main.py');
      await fs.writeFile(filePath, content, 'utf-8');

      await runPipedChild(io, sessionId, 'python3', ['-u', filePath], dir);
    } else if (language === 'java') {
      if (!tempDir) throw new Error('Temp dir not created');
      const dir = tempDir;
      const filePath = path.join(dir, 'Main.java');
      await fs.writeFile(filePath, content, 'utf-8');

      await new Promise<void>((resolve, reject) => {
        exec(`javac Main.java`, { cwd: dir, timeout: 5000 }, (error: Error | null, _stdout: string, stderr: string) => {
          if (error) {
            emitExec(io, sessionId, WS_EVENTS.EXECUTE_TOKEN, { token: stderr || error.message, isError: true });
            reject(new Error('Compilation Failed'));
          } else {
            resolve();
          }
        });
      });

      await runPipedChild(io, sessionId, 'java', ['Main'], dir);
    } else {
      emitExec(io, sessionId, WS_EVENTS.EXECUTE_TOKEN, { token: `Unsupported language: ${language}`, isError: true });
      emitExec(io, sessionId, WS_EVENTS.EXECUTE_COMPLETE, {});
    }
  } catch (err) {
    logger.error('Execution setup error', err);
    emitExec(io, sessionId, WS_EVENTS.EXECUTE_ERROR, { error: (err as Error).message });
    emitExec(io, sessionId, WS_EVENTS.EXECUTE_COMPLETE, {});
  } finally {
    if (tempDir) {
      try {
        await fs.rm(tempDir, { recursive: true, force: true });
      } catch (e) {
        logger.error(`Failed to cleanup temp dir ${tempDir}`, e);
      }
    }
  }
}

function isCurrentRun(sessionId: string, runId: string): boolean {
  return sessionActiveRunId.get(sessionId) === runId;
}

function runPipedChild(
  io: SocketServer,
  sessionId: string,
  command: string,
  args: string[],
  cwd: string
): Promise<void> {
  return new Promise((resolve) => {
    const proc = spawn(command, args, {
      cwd,
      env: { ...process.env, PYTHONUNBUFFERED: '1' } as NodeJS.ProcessEnv,
      stdio: ['pipe', 'pipe', 'pipe'],
    });

    const runId = randomUUID();
    sessionActiveRunId.set(sessionId, runId);
    activeChildProcs.set(runId, proc);

    emitExec(io, sessionId, WS_EVENTS.EXECUTE_STARTED, { runId });
    logger.debug(`Spawned ${command} run ${runId} for session ${sessionId}: ${args.join(' ')}`);

    let isFinished = false;
    let settled = false;

    const cleanup = () => {
      activeChildProcs.delete(runId);
      if (sessionActiveRunId.get(sessionId) === runId) {
        sessionActiveRunId.delete(sessionId);
      }
    };

    const releaseSuperseded = () => {
      if (settled) return;
      settled = true;
      isFinished = true;
      clearTimeout(timeoutPath);
      activeChildProcs.delete(runId);
      resolve();
    };

    const finish = () => {
      if (settled) return;
      settled = true;
      isFinished = true;
      clearTimeout(timeoutPath);
      cleanup();
      resolve();
    };

    const timeoutPath = setTimeout(() => {
      if (!isFinished && isCurrentRun(sessionId, runId)) {
        proc.kill('SIGKILL');
        emitExec(io, sessionId, WS_EVENTS.EXECUTE_TOKEN, {
          token: `\n[Execution Terminated: Timeout (${EXECUTION_TIMEOUT_MS / 1000}s)]`,
          isError: true,
        });
      }
    }, EXECUTION_TIMEOUT_MS);

    proc.stdout?.on('data', (chunk: Buffer) => {
      if (!isCurrentRun(sessionId, runId)) return;
      emitExec(io, sessionId, WS_EVENTS.EXECUTE_TOKEN, { token: chunk.toString('utf8') });
    });
    proc.stderr?.on('data', (chunk: Buffer) => {
      if (!isCurrentRun(sessionId, runId)) return;
      emitExec(io, sessionId, WS_EVENTS.EXECUTE_TOKEN, { token: chunk.toString('utf8'), isError: true });
    });

    proc.on('error', (err) => {
      if (!isCurrentRun(sessionId, runId)) {
        releaseSuperseded();
        return;
      }
      emitExec(io, sessionId, WS_EVENTS.EXECUTE_ERROR, { error: err.message });
      emitExec(io, sessionId, WS_EVENTS.EXECUTE_COMPLETE, {});
      finish();
    });

    proc.on('close', (code) => {
      if (settled) return;
      if (!isCurrentRun(sessionId, runId)) {
        releaseSuperseded();
        return;
      }
      emitExec(io, sessionId, WS_EVENTS.EXECUTE_TOKEN, { token: `\n[Process exited with code ${code}]` });
      emitExec(io, sessionId, WS_EVENTS.EXECUTE_COMPLETE, {});
      finish();
    });
  });
}

export function sendProcessInput(runId: string, input: string): boolean {
  const child = activeChildProcs.get(runId);
  if (!child?.stdin || child.stdin.destroyed) {
    logger.warn(`No active process for runId ${runId}. procs=${activeChildProcs.size}`);
    return false;
  }

  logger.debug(`stdin for run ${runId}: ${JSON.stringify(input)}`);
  child.stdin.write(Buffer.from(input, 'utf8'));
  return true;
}

export function cleanupProcess(sessionId: string): void {
  cancelGraceKill(sessionId);
  const runId = sessionActiveRunId.get(sessionId);
  if (runId) {
    const child = activeChildProcs.get(runId);
    if (child) {
      child.kill('SIGKILL');
      activeChildProcs.delete(runId);
    }
    sessionActiveRunId.delete(sessionId);
  }
}

	import { randomUUID } from 'crypto';
	import type { Server as SocketServer } from 'socket.io';
	import { exec, spawn, type ChildProcess } from 'child_process';
	import { promises as fs } from 'fs';
	import path from 'path';
	import os from 'os';
	import { WS_EVENTS } from '@asipilot/shared';
	import { logger } from '../utils/logger.js';
	import type { WSExecuteRequest } from '@asipilot/shared';
	import { getStableSessionId } from './session-id.js';

	/** Interactive runs may sit on input(); allow long thought + typing. */
	const EXECUTION_TIMEOUT_MS = 120_000;

	/** Child processes keyed by per-run id (client sends this with stdin). */
	const activeChildProcs = new Map<string, ChildProcess>();
	/** Latest run id for a browser tab session (cleared when run ends or new run starts). */
	const sessionActiveRunId = new Map<string, string>();

	const graceKillTimers = new Map<string, ReturnType<typeof setTimeout>>();

	export function cancelGraceKill(sessionId: string): void {
	const t = graceKillTimers.get(sessionId);
	if (t !== undefined) {
	clearTimeout(t);
	graceKillTimers.delete(sessionId);
	}
	}

	export function scheduleGraceKill(sessionId: string, delayMs = 15000): void {
	cancelGraceKill(sessionId);
	const t = setTimeout(() => {
	graceKillTimers.delete(sessionId);
	cleanupProcess(sessionId);
	logger.debug(`Grace-killed runner for session ${sessionId} after disconnect`);
	}, delayMs);
	graceKillTimers.set(sessionId, t);
	}

	/** Deliver to every socket whose handshake matches this tab session (rooms alone can fail behind some proxies). */
	function emitExec(io: SocketServer, sessionId: string, event: string, payload: object) {
	let n = 0;
	for (const socket of io.sockets.sockets.values()) {
	if (getStableSessionId(socket) === sessionId) {
	socket.emit(event, payload);
	n++;
	}
	}
	if (n === 0) {
	logger.warn(`exec ${event}: no socket matched session ${sessionId}; using room fallback`);
	io.to(sessionId).emit(event, payload);
	}
	}

	export async function handleCodeExecution(io: SocketServer, sessionId: string, data: WSExecuteRequest) {
	const { language, content } = data;
	let tempDir: string \| null = null;

	cancelGraceKill(sessionId);
	cleanupProcess(sessionId);

	try {
	tempDir = await fs.mkdtemp(path.join(os.tmpdir(), 'asipilot-exec-'));

	if (language === 'python') {
	if (!tempDir) throw new Error('Temp dir not created');
	const dir = tempDir;
	const filePath = path.join(dir, 'main.py');
	await fs.writeFile(filePath, content, 'utf-8');

	await runPipedChild(io, sessionId, 'python3', ['-u', filePath], dir);
	} else if (language === 'java') {
	if (!tempDir) throw new Error('Temp dir not created');
	const dir = tempDir;
	const filePath = path.join(dir, 'Main.java');
	await fs.writeFile(filePath, content, 'utf-8');

	await new Promise<void>((resolve, reject) => {
	exec(`javac Main.java`, { cwd: dir, timeout: 5000 }, (error: Error \| null, _stdout: string, stderr: string) => {
	if (error) {
	emitExec(io, sessionId, WS_EVENTS.EXECUTE_TOKEN, { token: stderr \|\| error.message, isError: true });
	reject(new Error('Compilation Failed'));
	} else {
	resolve();
	}
	});
	});

	await runPipedChild(io, sessionId, 'java', ['Main'], dir);
	} else {
	emitExec(io, sessionId, WS_EVENTS.EXECUTE_TOKEN, { token: `Unsupported language: ${language}`, isError: true });
	emitExec(io, sessionId, WS_EVENTS.EXECUTE_COMPLETE, {});
	}
	} catch (err) {
	logger.error('Execution setup error', err);
	emitExec(io, sessionId, WS_EVENTS.EXECUTE_ERROR, { error: (err as Error).message });
	emitExec(io, sessionId, WS_EVENTS.EXECUTE_COMPLETE, {});
	} finally {
	if (tempDir) {
	try {
	await fs.rm(tempDir, { recursive: true, force: true });
	} catch (e) {
	logger.error(`Failed to cleanup temp dir ${tempDir}`, e);
	}
	}
	}
	}

	function isCurrentRun(sessionId: string, runId: string): boolean {
	return sessionActiveRunId.get(sessionId) === runId;
	}

	function runPipedChild(
	io: SocketServer,
	sessionId: string,
	command: string,
	args: string[],
	cwd: string
	): Promise<void> {
	return new Promise((resolve) => {
	const proc = spawn(command, args, {
	cwd,
	env: { ...process.env, PYTHONUNBUFFERED: '1' } as NodeJS.ProcessEnv,
	stdio: ['pipe', 'pipe', 'pipe'],
	});

	const runId = randomUUID();
	sessionActiveRunId.set(sessionId, runId);
	activeChildProcs.set(runId, proc);

	emitExec(io, sessionId, WS_EVENTS.EXECUTE_STARTED, { runId });
	logger.debug(`Spawned ${command} run ${runId} for session ${sessionId}: ${args.join(' ')}`);

	let isFinished = false;
	let settled = false;

	const cleanup = () => {
	activeChildProcs.delete(runId);
	if (sessionActiveRunId.get(sessionId) === runId) {
	sessionActiveRunId.delete(sessionId);
	}
	};

	const releaseSuperseded = () => {
	if (settled) return;
	settled = true;
	isFinished = true;
	clearTimeout(timeoutPath);
	activeChildProcs.delete(runId);
	resolve();
	};

	const finish = () => {
	if (settled) return;
	settled = true;
	isFinished = true;
	clearTimeout(timeoutPath);
	cleanup();
	resolve();
	};

	const timeoutPath = setTimeout(() => {
	if (!isFinished && isCurrentRun(sessionId, runId)) {
	proc.kill('SIGKILL');
	emitExec(io, sessionId, WS_EVENTS.EXECUTE_TOKEN, {
	token: `\n[Execution Terminated: Timeout (${EXECUTION_TIMEOUT_MS / 1000}s)]`,
	isError: true,
	});
	}
	}, EXECUTION_TIMEOUT_MS);

	proc.stdout?.on('data', (chunk: Buffer) => {
	if (!isCurrentRun(sessionId, runId)) return;
	emitExec(io, sessionId, WS_EVENTS.EXECUTE_TOKEN, { token: chunk.toString('utf8') });
	});
	proc.stderr?.on('data', (chunk: Buffer) => {
	if (!isCurrentRun(sessionId, runId)) return;
	emitExec(io, sessionId, WS_EVENTS.EXECUTE_TOKEN, { token: chunk.toString('utf8'), isError: true });
	});

	proc.on('error', (err) => {
	if (!isCurrentRun(sessionId, runId)) {
	releaseSuperseded();
	return;
	}
	emitExec(io, sessionId, WS_EVENTS.EXECUTE_ERROR, { error: err.message });
	emitExec(io, sessionId, WS_EVENTS.EXECUTE_COMPLETE, {});
	finish();
	});

	proc.on('close', (code) => {
	if (settled) return;
	if (!isCurrentRun(sessionId, runId)) {
	releaseSuperseded();
	return;
	}
	emitExec(io, sessionId, WS_EVENTS.EXECUTE_TOKEN, { token: `\n[Process exited with code ${code}]` });
	emitExec(io, sessionId, WS_EVENTS.EXECUTE_COMPLETE, {});
	finish();
	});
	});
	}

	export function sendProcessInput(runId: string, input: string): boolean {
	const child = activeChildProcs.get(runId);
	if (!child?.stdin \|\| child.stdin.destroyed) {
	logger.warn(`No active process for runId ${runId}. procs=${activeChildProcs.size}`);
	return false;
	}

	logger.debug(`stdin for run ${runId}: ${JSON.stringify(input)}`);
	child.stdin.write(Buffer.from(input, 'utf8'));
	return true;
	}

	export function cleanupProcess(sessionId: string): void {
	cancelGraceKill(sessionId);
	const runId = sessionActiveRunId.get(sessionId);
	if (runId) {
	const child = activeChildProcs.get(runId);
	if (child) {
	child.kill('SIGKILL');
	activeChildProcs.delete(runId);
	}
	sessionActiveRunId.delete(sessionId);
	}
	}