628 lines
23 KiB
TypeScript
628 lines
23 KiB
TypeScript
/**
|
|
* Front server for the UI: fast file downloads + reverse proxy + Next.js supervisor.
|
|
*
|
|
* Owns the single public port. GET/HEAD /api/files/* are served directly with
|
|
* raw Node streams (full HTTP Range support, no Next.js per-request overhead);
|
|
* everything else — including the dev-mode HMR WebSocket — is piped unmodified
|
|
* to a Next.js server this process spawns on an ephemeral, loopback-only port.
|
|
* The internal port is never fixed (no collisions) and is not reachable from
|
|
* the network, so the externally visible surface is identical to running
|
|
* `next start` alone: one port.
|
|
*
|
|
* With AI_TOOLKIT_FILE_SERVER_WORKERS > 1 (default: min(4, cores) in start
|
|
* mode, 1 in dev mode), the front server runs as that many cluster workers
|
|
* sharing the public socket, so file serving is parallel across CPU cores and
|
|
* a burst of thumbnail requests never queues behind a multi-GB download.
|
|
*
|
|
* This process is optional. The equivalent (slower) route at
|
|
* src/app/api/files/[...filePath]/route.ts still exists, so deployments that
|
|
* run the worker and `next start` directly keep working exactly as before.
|
|
*
|
|
* Usage:
|
|
* node dist/cron/fileServer.js start --port 8675
|
|
* ts-node-dev cron/fileServer.ts dev --port 3000
|
|
*/
|
|
import cluster from 'cluster';
|
|
import { spawn } from 'child_process';
|
|
import http from 'http';
|
|
import net from 'net';
|
|
import fs from 'fs';
|
|
import os from 'os';
|
|
import path from 'path';
|
|
import { pipeline } from 'stream';
|
|
import prisma from './prisma';
|
|
import { defaultDatasetsFolder, defaultTrainFolder, defaultDataRoot, TOOLKIT_ROOT } from './paths';
|
|
|
|
const isDev = process.argv.includes('dev');
|
|
|
|
function argValue(name: string, fallback: number): number {
|
|
const i = process.argv.indexOf(name);
|
|
if (i !== -1 && process.argv[i + 1]) {
|
|
const v = parseInt(process.argv[i + 1], 10);
|
|
if (!isNaN(v)) return v;
|
|
}
|
|
return fallback;
|
|
}
|
|
|
|
const PUBLIC_PORT = argValue('--port', isDev ? 3000 : 8675);
|
|
const UPSTREAM_HOST = '127.0.0.1';
|
|
|
|
const numWorkers = (() => {
|
|
const env = parseInt(process.env.AI_TOOLKIT_FILE_SERVER_WORKERS || '', 10);
|
|
if (!isNaN(env) && env > 0) return env;
|
|
return isDev ? 1 : Math.min(4, os.cpus().length);
|
|
})();
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// File serving (same security model as the Next.js route). Settings can be
|
|
// changed at runtime from the UI in the Next.js process, so cache briefly
|
|
// instead of forever.
|
|
// ---------------------------------------------------------------------------
|
|
type Roots = { datasets: string; training: string; data: string };
|
|
let rootsCache: { roots: Roots; ts: number } | null = null;
|
|
|
|
async function getRoots(forceFresh = false): Promise<Roots> {
|
|
if (!forceFresh && rootsCache && Date.now() - rootsCache.ts < 10_000) {
|
|
return rootsCache.roots;
|
|
}
|
|
const rows = await prisma.settings.findMany({
|
|
where: { key: { in: ['DATASETS_FOLDER', 'TRAINING_FOLDER', 'DATA_ROOT'] } },
|
|
});
|
|
const fromRow = (key: string, fallback: string) => {
|
|
const row = rows.find(r => r.key === key);
|
|
// path.resolve strips trailing slashes; a root stored as "/mnt/foo/" would
|
|
// otherwise make the `root + path.sep` prefix check fail on every file.
|
|
return path.resolve(row?.value && row.value !== '' ? row.value : fallback);
|
|
};
|
|
const roots: Roots = {
|
|
datasets: fromRow('DATASETS_FOLDER', defaultDatasetsFolder),
|
|
training: fromRow('TRAINING_FOLDER', defaultTrainFolder),
|
|
data: fromRow('DATA_ROOT', defaultDataRoot),
|
|
};
|
|
rootsCache = { roots, ts: Date.now() };
|
|
return roots;
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Thumbnail generation for ?thumb=1 requests whose thumb doesn't exist yet.
|
|
// Output matches the Python generator (SampleConfig._generate_thumbnail in
|
|
// toolkit/config_modules.py): 300x300 center-cropped q90 jpg written
|
|
// atomically into the sibling .thumbs folder as <name>.<ext>.jpg.
|
|
// ---------------------------------------------------------------------------
|
|
const THUMB_SIZE = 300;
|
|
const IMAGE_THUMB_EXTS = new Set(['.jpg', '.jpeg', '.png', '.webp', '.gif', '.bmp']);
|
|
const VIDEO_THUMB_EXTS = new Set(['.mp4', '.avi', '.mov', '.mkv', '.wmv', '.m4v', '.flv']);
|
|
|
|
// sharp is not a direct dependency; Next.js vendors it (for next/image), so
|
|
// resolve it out of next's node_modules tree. If that ever fails, image
|
|
// thumbs just fall back to serving the original file.
|
|
const sharp: any = (() => {
|
|
try {
|
|
return require(require.resolve('sharp', { paths: [path.dirname(require.resolve('next/package.json'))] }));
|
|
} catch {
|
|
return null;
|
|
}
|
|
})();
|
|
|
|
// The manager provisions a portable FFmpeg at <repo>/.ffmpeg (see
|
|
// manager/ffmpeg.py) — prefer it over whatever is on PATH. Its Linux build is
|
|
// a shared one, so spawning it directly (i.e. not via `manager launch`, which
|
|
// sets this up itself) needs .ffmpeg/lib on LD_LIBRARY_PATH.
|
|
const localFfmpegExe = path.join(TOOLKIT_ROOT, '.ffmpeg', 'bin', process.platform === 'win32' ? 'ffmpeg.exe' : 'ffmpeg');
|
|
const ffmpegExe = fs.existsSync(localFfmpegExe) ? localFfmpegExe : 'ffmpeg';
|
|
const ffmpegEnv: NodeJS.ProcessEnv = (() => {
|
|
const libDir = path.join(TOOLKIT_ROOT, '.ffmpeg', 'lib');
|
|
if (ffmpegExe === 'ffmpeg' || process.platform !== 'linux' || !fs.existsSync(libDir)) return process.env;
|
|
const prior = process.env.LD_LIBRARY_PATH;
|
|
return { ...process.env, LD_LIBRARY_PATH: prior ? `${libDir}${path.delimiter}${prior}` : libDir };
|
|
})();
|
|
|
|
let ffmpegMissing = false;
|
|
// A generation attempt that failed is memoized (keyed on source mtime, so a
|
|
// re-written file retries) — the gallery re-requests thumbs constantly and
|
|
// must not re-run ffmpeg/sharp against a broken file on every poll.
|
|
const failedThumbs = new Set<string>();
|
|
const inFlightThumbs = new Map<string, Promise<boolean>>();
|
|
|
|
// A gallery burst can request hundreds of missing thumbs at once; cap how
|
|
// many decodes/ffmpeg spawns run concurrently per worker.
|
|
const MAX_THUMB_GEN = 4;
|
|
let activeThumbGen = 0;
|
|
const thumbGenQueue: (() => void)[] = [];
|
|
async function withThumbGenSlot<T>(fn: () => Promise<T>): Promise<T> {
|
|
if (activeThumbGen >= MAX_THUMB_GEN) {
|
|
await new Promise<void>(r => thumbGenQueue.push(r));
|
|
}
|
|
activeThumbGen++;
|
|
try {
|
|
return await fn();
|
|
} finally {
|
|
activeThumbGen--;
|
|
thumbGenQueue.shift()?.();
|
|
}
|
|
}
|
|
|
|
async function generateThumb(sourcePath: string, thumbPath: string): Promise<boolean> {
|
|
const ext = path.extname(sourcePath).toLowerCase();
|
|
const isImage = IMAGE_THUMB_EXTS.has(ext);
|
|
const isVideo = VIDEO_THUMB_EXTS.has(ext);
|
|
if ((isImage && !sharp) || (isVideo && ffmpegMissing) || (!isImage && !isVideo)) return false;
|
|
await fs.promises.mkdir(path.dirname(thumbPath), { recursive: true });
|
|
// Write to a per-process tmp name, then atomically rename into place (same
|
|
// as the Python generator) so a concurrent request never reads a partial
|
|
// thumb. The .jpg suffix is required for ffmpeg's output format detection.
|
|
const tmpPath = `${thumbPath}.${process.pid}.tmp.jpg`;
|
|
try {
|
|
if (isImage) {
|
|
// sharp opens animated formats on the first frame by default
|
|
await sharp(sourcePath)
|
|
.resize(THUMB_SIZE, THUMB_SIZE, { fit: 'cover' })
|
|
.jpeg({ quality: 90 })
|
|
.toFile(tmpPath);
|
|
} else {
|
|
await new Promise<void>((resolve, reject) => {
|
|
const child = spawn(
|
|
ffmpegExe,
|
|
[
|
|
'-y',
|
|
'-loglevel', 'error',
|
|
'-i', sourcePath,
|
|
'-frames:v', '1',
|
|
'-vf', `crop='min(iw,ih)':'min(iw,ih)',scale=${THUMB_SIZE}:${THUMB_SIZE}`,
|
|
'-q:v', '2',
|
|
tmpPath,
|
|
],
|
|
{ stdio: ['ignore', 'ignore', 'pipe'], env: ffmpegEnv },
|
|
);
|
|
let stderr = '';
|
|
child.stderr!.on('data', chunk => (stderr += chunk.toString()));
|
|
const timer = setTimeout(() => child.kill('SIGKILL'), 30_000);
|
|
child.on('error', (err: NodeJS.ErrnoException) => {
|
|
clearTimeout(timer);
|
|
if (err.code === 'ENOENT') ffmpegMissing = true;
|
|
reject(err);
|
|
});
|
|
child.on('exit', code => {
|
|
clearTimeout(timer);
|
|
code === 0 ? resolve() : reject(new Error(`ffmpeg exited with ${code}: ${stderr.trim()}`));
|
|
});
|
|
});
|
|
}
|
|
await fs.promises.rename(tmpPath, thumbPath);
|
|
return true;
|
|
} catch (err) {
|
|
await fs.promises.unlink(tmpPath).catch(() => { });
|
|
throw err;
|
|
}
|
|
}
|
|
|
|
function ensureThumb(sourcePath: string, thumbPath: string, sourceMtimeMs: number): Promise<boolean> {
|
|
const failKey = `${thumbPath}:${sourceMtimeMs}`;
|
|
if (failedThumbs.has(failKey)) return Promise.resolve(false);
|
|
let pending = inFlightThumbs.get(thumbPath);
|
|
if (!pending) {
|
|
pending = withThumbGenSlot(() => generateThumb(sourcePath, thumbPath))
|
|
.catch(err => {
|
|
console.warn(`Failed to generate thumbnail for ${sourcePath}: ${err?.message || err}`);
|
|
return false;
|
|
})
|
|
.then(ok => {
|
|
if (!ok) failedThumbs.add(failKey);
|
|
return ok;
|
|
})
|
|
.finally(() => inFlightThumbs.delete(thumbPath));
|
|
inFlightThumbs.set(thumbPath, pending);
|
|
}
|
|
return pending;
|
|
}
|
|
|
|
const contentTypeMap: { [key: string]: string } = {
|
|
'.jpg': 'image/jpeg',
|
|
'.jpeg': 'image/jpeg',
|
|
'.png': 'image/png',
|
|
'.gif': 'image/gif',
|
|
'.webp': 'image/webp',
|
|
'.svg': 'image/svg+xml',
|
|
'.bmp': 'image/bmp',
|
|
'.safetensors': 'application/octet-stream',
|
|
'.zip': 'application/zip',
|
|
// Videos
|
|
'.mp4': 'video/mp4',
|
|
'.avi': 'video/x-msvideo',
|
|
'.mov': 'video/quicktime',
|
|
'.mkv': 'video/x-matroska',
|
|
'.wmv': 'video/x-ms-wmv',
|
|
'.m4v': 'video/x-m4v',
|
|
'.flv': 'video/x-flv',
|
|
// Audio
|
|
'.mp3': 'audio/mpeg',
|
|
'.wav': 'audio/wav',
|
|
'.flac': 'audio/flac',
|
|
'.ogg': 'audio/ogg',
|
|
};
|
|
|
|
async function serveFile(req: http.IncomingMessage, res: http.ServerResponse, prefix: string): Promise<void> {
|
|
try {
|
|
// /api/img/ serves media inline with ETag revalidation; /api/files/ is a
|
|
// forced download. Both mirror their Next.js route's exact semantics.
|
|
const isImg = prefix === '/api/img/';
|
|
const urlPath = (req.url || '').split('?')[0];
|
|
const rest = urlPath.slice(prefix.length);
|
|
// Decode per URL segment so both forms resolve to the same file:
|
|
// /api/files/%2Fmnt%2Fout%2Fjob%2Ffile.safetensors (legacy: whole path in one segment)
|
|
// /api/files/%2Fmnt%2Fout%2Fjob/file.safetensors (folder / filename — what the UI emits;
|
|
// gives wget & co. the real filename)
|
|
// Windows folders arrive as `C%3A%5Cout%5Cjob` and decode to `C:\out\job`;
|
|
// path.resolve below normalizes the mixed `\`/`/` separators. Same logic as
|
|
// catchAllToFilePath in src/server/catchAllPath.ts for the Next.js routes.
|
|
const decodedFilePath = rest
|
|
.split('/')
|
|
.map(decodeURIComponent)
|
|
.join('/');
|
|
|
|
let resolvedFilePath = path.resolve(decodedFilePath);
|
|
const checkAllowed = (roots: Roots) => {
|
|
const dirs = isImg ? [roots.datasets, roots.training, roots.data] : [roots.datasets, roots.training];
|
|
return {
|
|
dirs,
|
|
allowed: dirs.some(
|
|
allowedDir => resolvedFilePath === allowedDir || resolvedFilePath.startsWith(allowedDir + path.sep),
|
|
),
|
|
};
|
|
};
|
|
let { dirs: allowedDirs, allowed: isAllowed } = checkAllowed(await getRoots());
|
|
if (!isAllowed) {
|
|
// The cached roots may be stale — settings can be changed from the UI at
|
|
// any moment. Re-fetch before denying so a just-updated path never 403s.
|
|
({ dirs: allowedDirs, allowed: isAllowed } = checkAllowed(await getRoots(true)));
|
|
}
|
|
if (!isAllowed) {
|
|
console.warn(`Access denied: ${resolvedFilePath} not in ${allowedDirs.join(', ')}`);
|
|
res.writeHead(403);
|
|
res.end('Access denied');
|
|
return;
|
|
}
|
|
|
|
// ?thumb=1 serves the 300x300 jpg from the sibling .thumbs folder
|
|
// (<name>.<ext>.jpg), generating and saving it on the fly when missing.
|
|
// Falls through to the full file only if generation isn't possible
|
|
// (unsupported format, no ffmpeg, corrupt file).
|
|
if (isImg && new URL(req.url || '', 'http://localhost').searchParams.has('thumb')) {
|
|
const thumbPath = path.join(path.dirname(resolvedFilePath), '.thumbs', path.basename(resolvedFilePath) + '.jpg');
|
|
let thumbStat = await fs.promises.stat(thumbPath).catch(() => null);
|
|
if (!(thumbStat && thumbStat.isFile())) {
|
|
const srcStat = await fs.promises.stat(resolvedFilePath).catch(() => null);
|
|
if (srcStat && srcStat.isFile() && (await ensureThumb(resolvedFilePath, thumbPath, srcStat.mtimeMs))) {
|
|
thumbStat = await fs.promises.stat(thumbPath).catch(() => null);
|
|
}
|
|
}
|
|
if (thumbStat && thumbStat.isFile()) {
|
|
resolvedFilePath = thumbPath;
|
|
}
|
|
}
|
|
|
|
let stat: fs.Stats;
|
|
try {
|
|
stat = await fs.promises.stat(resolvedFilePath);
|
|
} catch {
|
|
res.writeHead(404);
|
|
res.end('File not found');
|
|
return;
|
|
}
|
|
if (!stat.isFile()) {
|
|
res.writeHead(400);
|
|
res.end('Not a file');
|
|
return;
|
|
}
|
|
|
|
const filename = path.basename(resolvedFilePath);
|
|
const ext = path.extname(resolvedFilePath).toLowerCase();
|
|
const contentType = contentTypeMap[ext] || 'application/octet-stream';
|
|
|
|
const headers: http.OutgoingHttpHeaders = {
|
|
'Content-Type': contentType,
|
|
'Accept-Ranges': 'bytes',
|
|
};
|
|
if (isImg) {
|
|
// Weak ETag from inode/size/mtime — cheap and stable enough for
|
|
// revalidation of thumbnails/videos the UI requests over and over.
|
|
const etag = `W/"${stat.ino.toString(36)}-${stat.size.toString(36)}-${stat.mtimeMs.toString(36)}"`;
|
|
headers['Cache-Control'] = 'public, max-age=86400, immutable';
|
|
headers['ETag'] = etag;
|
|
if (req.headers['if-none-match'] === etag) {
|
|
res.writeHead(304, { ETag: etag, 'Cache-Control': headers['Cache-Control'] });
|
|
res.end();
|
|
return;
|
|
}
|
|
} else {
|
|
headers['Cache-Control'] = 'public, max-age=86400';
|
|
headers['Content-Disposition'] = `attachment; filename="${encodeURIComponent(filename)}"`;
|
|
headers['X-Content-Type-Options'] = 'nosniff';
|
|
}
|
|
|
|
const range = req.headers.range;
|
|
let start = 0;
|
|
let end = stat.size - 1;
|
|
let status = 200;
|
|
|
|
if (range) {
|
|
// Open-ended ranges (`bytes=0-`) serve to EOF — capping them forces
|
|
// clients into serial re-requests and RTT-bound throughput.
|
|
const parts = range.replace(/bytes=/, '').split('-');
|
|
start = parseInt(parts[0], 10);
|
|
end = parts[1] ? parseInt(parts[1], 10) : stat.size - 1;
|
|
if (isNaN(start) || isNaN(end) || start > end || start >= stat.size) {
|
|
res.writeHead(416, { 'Content-Range': `bytes */${stat.size}` });
|
|
res.end();
|
|
return;
|
|
}
|
|
status = 206;
|
|
headers['Content-Range'] = `bytes ${start}-${end}/${stat.size}`;
|
|
}
|
|
|
|
headers['Content-Length'] = String(end - start + 1);
|
|
res.writeHead(status, headers);
|
|
if (req.method === 'HEAD') {
|
|
res.end();
|
|
return;
|
|
}
|
|
if (res.destroyed) {
|
|
// Client disconnected while we were stat-ing; piping into a destroyed
|
|
// stream throws.
|
|
return;
|
|
}
|
|
|
|
const fileStream = fs.createReadStream(resolvedFilePath, {
|
|
start,
|
|
end,
|
|
highWaterMark: 4 * 1024 * 1024, // large buffer to minimize per-chunk overhead
|
|
});
|
|
pipeline(fileStream, res, err => {
|
|
// Client disconnects mid-download land here; just make sure the socket
|
|
// is gone.
|
|
if (err) res.destroy();
|
|
});
|
|
} catch (error) {
|
|
console.error('Error serving file:', error);
|
|
if (!res.headersSent) res.writeHead(500);
|
|
res.end();
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Passthrough proxy to the Next.js server
|
|
// ---------------------------------------------------------------------------
|
|
const upstreamAgent = new http.Agent({ keepAlive: true });
|
|
|
|
// Hop-by-hop headers must not be forwarded: Node de-chunks bodies on read and
|
|
// re-frames them on write, so forwarding e.g. `transfer-encoding` would
|
|
// corrupt the stream.
|
|
const HOP_BY_HOP = [
|
|
'connection',
|
|
'keep-alive',
|
|
'proxy-authenticate',
|
|
'proxy-authorization',
|
|
'te',
|
|
'trailer',
|
|
'transfer-encoding',
|
|
'upgrade',
|
|
];
|
|
|
|
function stripHopByHop(headers: http.IncomingHttpHeaders): http.IncomingHttpHeaders {
|
|
const out = { ...headers };
|
|
for (const h of HOP_BY_HOP) delete out[h];
|
|
return out;
|
|
}
|
|
|
|
function proxy(req: http.IncomingMessage, res: http.ServerResponse, upstreamPort: number, attempt = 0): void {
|
|
const bodyless = req.method === 'GET' || req.method === 'HEAD';
|
|
const upstreamReq = http.request(
|
|
{
|
|
host: UPSTREAM_HOST,
|
|
port: upstreamPort,
|
|
path: req.url,
|
|
method: req.method,
|
|
headers: stripHopByHop(req.headers),
|
|
agent: upstreamAgent,
|
|
},
|
|
upstreamRes => {
|
|
if (res.destroyed) {
|
|
// Client disconnected while Next.js was working; drain and drop.
|
|
upstreamRes.resume();
|
|
return;
|
|
}
|
|
res.writeHead(upstreamRes.statusCode || 502, stripHopByHop(upstreamRes.headers));
|
|
pipeline(upstreamRes, res, () => { });
|
|
},
|
|
);
|
|
// If the client goes away, stop the upstream work too. No-op after normal
|
|
// completion.
|
|
const abortUpstream = () => upstreamReq.destroy();
|
|
res.once('close', abortUpstream);
|
|
upstreamReq.on('error', (err: NodeJS.ErrnoException) => {
|
|
res.removeListener('close', abortUpstream);
|
|
// Next.js is still booting. Nothing was sent on a refused connection, so
|
|
// bodyless requests are safe to retry (request bodies can't be replayed).
|
|
if (err.code === 'ECONNREFUSED' && bodyless && attempt < 120 && !res.destroyed) {
|
|
setTimeout(() => proxy(req, res, upstreamPort, attempt + 1), 250);
|
|
return;
|
|
}
|
|
// Keep-alive reuse race: Next.js closes pooled sockets after 5s idle, and
|
|
// the UI polls on ~5s intervals, so a reused socket can die the instant we
|
|
// write to it (ECONNRESET/EPIPE). No response bytes exist yet, so retrying
|
|
// immediately is safe; each retry drains one stale socket from the pool
|
|
// until a live or fresh connection is used.
|
|
if (
|
|
bodyless &&
|
|
upstreamReq.reusedSocket &&
|
|
(err.code === 'ECONNRESET' || err.code === 'EPIPE') &&
|
|
!res.headersSent &&
|
|
!res.destroyed &&
|
|
attempt < 120
|
|
) {
|
|
proxy(req, res, upstreamPort, attempt + 1);
|
|
return;
|
|
}
|
|
if (res.destroyed) return;
|
|
if (!res.headersSent) res.writeHead(502);
|
|
res.end('UI server unavailable');
|
|
});
|
|
if (bodyless) {
|
|
upstreamReq.end();
|
|
} else {
|
|
pipeline(req, upstreamReq, () => { });
|
|
}
|
|
}
|
|
|
|
function startServer(publicPort: number, upstreamPort: number): void {
|
|
const server = http.createServer((req, res) => {
|
|
const urlPath = (req.url || '').split('?')[0];
|
|
if (req.method === 'GET' || req.method === 'HEAD') {
|
|
const prefix = ['/api/files/', '/api/img/'].find(p => urlPath.startsWith(p));
|
|
if (prefix) {
|
|
serveFile(req, res, prefix);
|
|
return;
|
|
}
|
|
}
|
|
proxy(req, res, upstreamPort);
|
|
});
|
|
|
|
// WebSocket upgrades (dev HMR) are forwarded as raw TCP with the original
|
|
// request head replayed, so the hop-by-hop upgrade headers pass through
|
|
// intact.
|
|
server.on('upgrade', (req, socket, head) => {
|
|
const upstream = net.connect(upstreamPort, UPSTREAM_HOST, () => {
|
|
if (socket.destroyed) {
|
|
upstream.destroy();
|
|
return;
|
|
}
|
|
let raw = `${req.method} ${req.url} HTTP/1.1\r\n`;
|
|
for (let i = 0; i < req.rawHeaders.length; i += 2) {
|
|
raw += `${req.rawHeaders[i]}: ${req.rawHeaders[i + 1]}\r\n`;
|
|
}
|
|
raw += '\r\n';
|
|
upstream.write(raw);
|
|
if (head && head.length) upstream.write(head);
|
|
upstream.pipe(socket);
|
|
socket.pipe(upstream);
|
|
});
|
|
upstream.on('error', () => socket.destroy());
|
|
socket.on('error', () => upstream.destroy());
|
|
});
|
|
|
|
// Uploads can be huge (100gb body limit upstream); the default 5-minute
|
|
// request timeout would kill them mid-transfer.
|
|
server.requestTimeout = 0;
|
|
|
|
server.listen(publicPort);
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Supervisor: spawn Next.js on an ephemeral loopback port, then serve
|
|
// ---------------------------------------------------------------------------
|
|
function getFreePort(): Promise<number> {
|
|
return new Promise((resolve, reject) => {
|
|
const probe = net.createServer();
|
|
probe.on('error', reject);
|
|
probe.listen(0, UPSTREAM_HOST, () => {
|
|
const port = (probe.address() as net.AddressInfo).port;
|
|
probe.close(err => (err ? reject(err) : resolve(port)));
|
|
});
|
|
});
|
|
}
|
|
|
|
async function primaryMain(): Promise<void> {
|
|
const upstreamPort = await getFreePort();
|
|
|
|
const nextBin = require.resolve('next/dist/bin/next');
|
|
const nextArgs = isDev ? ['dev', '--turbopack'] : ['start'];
|
|
nextArgs.push('--port', String(upstreamPort), '--hostname', UPSTREAM_HOST);
|
|
|
|
const nextChild = spawn(process.execPath, [nextBin, ...nextArgs], {
|
|
stdio: ['inherit', 'pipe', 'pipe'],
|
|
cwd: process.cwd(),
|
|
});
|
|
|
|
// Next.js prints its banner with the internal loopback port ("- Local:
|
|
// http://127.0.0.1:39865"). Users who follow that link get the
|
|
// non-accelerated server (or nothing, remotely). Rewrite every mention of
|
|
// the internal address to the real public port before it reaches the
|
|
// console.
|
|
const internalAddr = new RegExp(`(https?://)?(127\\.0\\.0\\.1|localhost):${upstreamPort}`, 'g');
|
|
const rewritePipe = (from: NodeJS.ReadableStream, to: NodeJS.WritableStream) => {
|
|
let buffered = '';
|
|
from.on('data', (chunk: Buffer) => {
|
|
buffered += chunk.toString();
|
|
const lines = buffered.split('\n');
|
|
buffered = lines.pop() || '';
|
|
for (const line of lines) {
|
|
to.write(line.replace(internalAddr, `http://localhost:${PUBLIC_PORT}`) + '\n');
|
|
}
|
|
});
|
|
from.on('end', () => {
|
|
if (buffered) to.write(buffered.replace(internalAddr, `http://localhost:${PUBLIC_PORT}`) + '\n');
|
|
});
|
|
};
|
|
rewritePipe(nextChild.stdout!, process.stdout);
|
|
rewritePipe(nextChild.stderr!, process.stderr);
|
|
|
|
let shuttingDown = false;
|
|
const shutdown = (code: number) => {
|
|
if (shuttingDown) return;
|
|
shuttingDown = true;
|
|
nextChild.kill('SIGTERM');
|
|
for (const worker of Object.values(cluster.workers || {})) {
|
|
worker?.kill();
|
|
}
|
|
process.exit(code);
|
|
};
|
|
|
|
nextChild.on('exit', code => {
|
|
if (!shuttingDown) {
|
|
console.error(`Next.js exited with code ${code}, shutting down file server`);
|
|
shutdown(code ?? 1);
|
|
}
|
|
});
|
|
process.on('SIGINT', () => shutdown(0));
|
|
process.on('SIGTERM', () => shutdown(0));
|
|
// Never leave an orphaned Next.js child behind if we die unexpectedly.
|
|
process.on('uncaughtException', err => {
|
|
console.error('File server uncaught exception:', err);
|
|
shutdown(1);
|
|
});
|
|
|
|
console.log(`AI Toolkit UI: http://localhost:${PUBLIC_PORT} (${numWorkers} file server worker${numWorkers === 1 ? '' : 's'})`);
|
|
|
|
if (numWorkers <= 1) {
|
|
startServer(PUBLIC_PORT, upstreamPort);
|
|
return;
|
|
}
|
|
|
|
// AI_TOOLKIT_QUIET_PATHS keeps each forked worker from re-printing the
|
|
// TOOLKIT_ROOT line from paths.ts.
|
|
const workerEnv = {
|
|
UPSTREAM_PORT: String(upstreamPort),
|
|
PUBLIC_PORT: String(PUBLIC_PORT),
|
|
AI_TOOLKIT_QUIET_PATHS: '1',
|
|
};
|
|
for (let i = 0; i < numWorkers; i++) {
|
|
cluster.fork(workerEnv);
|
|
}
|
|
cluster.on('exit', worker => {
|
|
if (!shuttingDown) {
|
|
console.warn(`File server worker ${worker.id} died, restarting`);
|
|
cluster.fork(workerEnv);
|
|
}
|
|
});
|
|
}
|
|
|
|
if (cluster.isPrimary) {
|
|
primaryMain().catch(err => {
|
|
console.error('File server failed to start:', err);
|
|
process.exit(1);
|
|
});
|
|
} else {
|
|
process.on('disconnect', () => process.exit(0));
|
|
startServer(parseInt(process.env.PUBLIC_PORT!, 10), parseInt(process.env.UPSTREAM_PORT!, 10));
|
|
}
|