@needle-tools/engine/plugins/vite/dependencies.js

Needle Engine is a web-based runtime for 3D apps. It runs on your machine for development with great integrations into editors like Unity or Blender - and can be deployed onto any device! It is flexible, extensible and networking and XR are built-in.

import path from 'path';
import { existsSync, readFileSync, readdirSync } from 'fs';
import { tryGetNeedleEngineVersion, tryGetPackageVersion } from '../common/version.js';
import { needleLog } from './logging.js';

/**
 * @type {string[]}
 */
export const preloadScriptPaths = [];


/**
 * Returns true when \`@needle-tools/engine\` is installed as a local package
 * (i.e. it has its own nested `node_modules`). Vite's optimiser must skip such
 * packages, otherwise it tries to pre-bundle source that was never meant to be
 * pre-bundled and fails at dev-server start.
 *
 * @param {string} [root] – project root; defaults to `process.cwd()`.
 * @returns {boolean}
 */
export function isLocalNeedleEngineInstalled(root = process.cwd()) {
    const lockPath = path.resolve(root, 'node_modules', '@needle-tools/engine', 'node_modules', '.package-lock.json');
    return existsSync(lockPath);
}

/**
 * @param {"build" | "serve"} command
 * @param {import('../types/needleConfig').needleMeta | null | undefined} config
 * @param {import('../types').userSettings} userSettings
 * @returns {import('vite').Plugin[]}
 */
export function needleDependencies(command, config, userSettings) {

    /**
     * @type {import('vite').Plugin}
     */
    return [
        {
            name: 'needle:dependencies',
            enforce: 'pre',
            /**
             * @param {import('vite').UserConfig} config
             */
            config: (config, env) => {
                handleOptimizeDeps(config);
                handleManualChunks(config);
                // Ensure all imports of three and needle-engine resolve to the same copy
                // This prevents duplicates when packages (e.g. sample scripts) have their own node_modules
                if (!config.resolve) config.resolve = {};
                if (!config.resolve.dedupe) config.resolve.dedupe = [];
                for (const pkg of ['three', '@needle-tools/engine']) {
                    if (!config.resolve.dedupe.includes(pkg)) {
                        config.resolve.dedupe.push(pkg);
                    }
                }
                // Return server.fs config as a partial so Vite deep-merges it correctly.
                // Mutating config.server directly is not reliable for nested objects.
                if (isLocalNeedleEngineInstalled()) {
                    needleLog("needle-dependencies", 'Detected local @needle-tools/engine package → disabling server.fs.strict to allow symlinked worker imports');
                    return { server: { fs: { strict: false } } };
                }
            },
        },
    ]
}

const excludeDependencies = [
]

/**
 * Scans the engine source to find all `three/examples/jsm/...` module paths
 * that are actually imported. Returns a regex matching those paths so they
 * go into the core `three-examples` chunk. Anything else from three/examples
 * that the engine doesn't use goes into `three-examples.extras`.
 * @returns {RegExp}
 */
function buildThreeExamplesCoreTest() {
    /** @type {Set<string>} */
    const modulePaths = new Set();
    // Find the engine source directory relative to this plugin file
    const srcDir = path.resolve(import.meta.dirname, '../../src');
    if (!existsSync(srcDir)) {
        needleLog('needle-dependencies', `Engine src not found at "${srcDir}" — using fallback regex`);
        return /three\/(?:examples\/jsm|addons)\/(?:loaders|libs)\//;
    }
    scanDir(srcDir);

    function scanDir(dir) {
        for (const entry of readdirSync(dir, { withFileTypes: true })) {
            if (entry.isDirectory()) {
                scanDir(path.join(dir, entry.name));
            }
            else if (entry.name.endsWith('.ts')) {
                try {
                    const content = readFileSync(path.join(dir, entry.name), 'utf8');
                    // Match: from "three/examples/jsm/..." or "three/addons/..." (skip type-only imports)
                    const re = /(?<!import\s+type\s[^;]*?)from\s+["']three\/(?:examples\/jsm|addons)\/([^"']+)["']/g;
                    let m;
                    while ((m = re.exec(content)) !== null) {
                        // Extract the directory part, e.g. "loaders/GLTFLoader.js" → "loaders"
                        const dir = m[1].split('/')[0];
                        modulePaths.add(dir);
                    }
                }
                catch { /* skip unreadable files */ }
            }
        }
    }

    if (modulePaths.size === 0) {
        return /three\/(?:examples\/jsm|addons)\/(?:loaders|libs)\//;
    }

    // Build regex: three/(examples/jsm|addons)/(loaders|libs|controls|...)/
    const dirs = [...modulePaths].sort();
    needleLog('needle-dependencies', `three-examples core dirs: ${dirs.join(', ')}`);
    const dirsPattern = dirs.join('|');
    const pattern = `three\\/(?:examples\\/jsm|addons)\\/(?:${dirsPattern})\\/`;
    return new RegExp(pattern);
}

/**
 * @param {import('vite').UserConfig} config
 */
function handleOptimizeDeps(config) {
    excludeDependencies.forEach(dep => {
        if (config.optimizeDeps?.include?.includes(dep)) {
            needleLog("needle-dependencies", `${dep} is included in the optimizeDeps.include array. This may cause issues with the worker import.`, "warn", { dimBody: false });
        }
        else {
            if (!config.optimizeDeps) {
                config.optimizeDeps = {};
            }
            if (!config.optimizeDeps.exclude) {
                config.optimizeDeps.exclude = [];
            }

            // This needs to be excluded from optimization because otherwise the worker import fails
            // three-mesh-bvh/src/workers/generateMeshBVH.worker.js?worker
            // same for gltf-progressive
            needleLog("needle-dependencies", `Adding ${dep} to the optimizeDeps.exclude array to support workers.`);
            config.optimizeDeps.exclude.push(dep);

            if (!config.server) config.server = {};
            if (!config.server.fs) config.server.fs = {};
            if (config.server.fs.strict === undefined) {
                // we need to disable strictness to allow importing the worker from three-mesh-bvh node_modules in our GenerateMeshBVHWorker.js file
                config.server.fs.strict = false;
            }
        }
    });
    // is needle engine a local package? exclude it from vite's optimizeDeps
    if (isLocalNeedleEngineInstalled()) {
        config.optimizeDeps ??= {};
        config.optimizeDeps.exclude ??= [];
        if (!config.optimizeDeps.include?.includes('@needle-tools/engine') &&
            !config.optimizeDeps.exclude.includes('@needle-tools/engine')) {
            config.optimizeDeps.exclude.push('@needle-tools/engine');
            needleLog("needle-dependencies", 'Detected local @needle-tools/engine package → will exclude it from optimization');
        }
        // When engine is excluded from optimizeDeps, three-mesh-bvh must also be excluded.
        // The BVH worker (generateMeshBVH.worker.js) uses bare imports like `import 'three'`
        // which only resolve correctly when served through Vite's dev server module system.
        // If three-mesh-bvh is pre-bundled, the worker URL points into the cache and bare
        // imports fail at runtime → "Unknown error. Please check the server console."
        if (!config.optimizeDeps.include?.includes('three-mesh-bvh') &&
            !config.optimizeDeps.exclude.includes('three-mesh-bvh')) {
            config.optimizeDeps.exclude.push('three-mesh-bvh');
            needleLog("needle-dependencies", 'Detected local @needle-tools/engine package → will also exclude three-mesh-bvh from optimization');
        }
    }
}

/**
 * @param {import('vite').UserConfig} config
 */
/** @type {RegExp} */
let threeExamplesCoreRegex;

function handleManualChunks(config) {
    if (!config.build) {
        config.build = {};
    }
    threeExamplesCoreRegex = buildThreeExamplesCoreTest();
    // Support both rolldownOptions (Vite 8+) and rollupOptions (Vite 7 and earlier)
    const optionsKey = 'rolldownOptions' in (config.build) ? 'rolldownOptions' : 'rollupOptions';
    if (!config.build[optionsKey]) {
        config.build[optionsKey] = {};
    }
    if (!config.build[optionsKey].output) {
        config.build[optionsKey].output = {};
    }

    const rollupOutput = config.build[optionsKey].output;

    // Detect Vite 8+ (Rolldown-based) which needs codeSplitting.groups instead of manualChunks.
    // In Rolldown, manualChunks gets converted to a single codeSplitting group with recursive
    // dependency inclusion, causing chunks like gltf-progressive to absorb three.js.
    let isVite8 = false;
    try {
        const vitePkgPath = path.resolve(process.cwd(), 'node_modules/vite/package.json');
        if (existsSync(vitePkgPath)) {
            isVite8 = parseInt(JSON.parse(readFileSync(vitePkgPath, 'utf8')).version) >= 8;
        }
    } catch { }
    if (isVite8) needleLog("needle-dependencies", `Vite 8+ detected`);

    if (Array.isArray(rollupOutput)) {
        // append the manualChunks function to the array
        if (process.env.DEBUG) needleLog("needle-dependencies", "registering manualChunks");
        rollupOutput.push({
            manualChunks: needleManualChunks
        })
    }
    else {
        // If preserveModules is true we can not modify the chunkFileNames
        let allowManualChunks = true;
        if ("manualChunks" in rollupOutput || "codeSplitting" in rollupOutput) {
            allowManualChunks = false;
            // if the user has already defined manualChunks/codeSplitting, we don't want to overwrite it
            needleLog("needle-dependencies", "manualChunks/codeSplitting already found in vite config - will not overwrite it");
        }
        else if (rollupOutput.preserveModules === true) {
            allowManualChunks = false;
            needleLog("needle-dependencies", "manualChunks can not be registered because preserveModules is true");
        }
        if (rollupOutput.inlineDynamicImports === true) {
            allowManualChunks = false;
            needleLog("needle-dependencies", "manualChunks can not be registered because inlineDynamicImports is true");
        }

        if (allowManualChunks) {
            if (process.env.DEBUG) needleLog("needle-dependencies", "registering manualChunks");
            if (isVite8) {
                // Vite 8+ uses Rolldown which converts manualChunks into a single codeSplitting group
                // with recursive dependency inclusion. This causes gltf-progressive to absorb three.js.
                // Use codeSplitting.groups directly with priorities so three.js gets its own chunk.
                // Priority determines which group claims a module first. With the default
                // includeDependenciesRecursively: true, claimed modules pull their deps into
                // the same chunk — UNLESS those deps were already claimed by a higher-priority
                // group. So three.js core (priority 100) must be higher than three-examples (90)
                // to prevent examples from absorbing the core lib via recursive deps.
                rollupOutput.codeSplitting = {
                    groups: [
                        { name: 'three', test: /\/three\/(?:build|src)\//, priority: 100 },
                        // three/examples modules used by the engine go into the core chunk.
                        // Anything else (user-installed three/examples) goes into extras.
                        { name: 'three-examples', test: threeExamplesCoreRegex, priority: 95 },
                        { name: 'three-examples.extras', test: /three\/(?:examples|addons)/, priority: 90 },
                        { name: 'three-mesh-bvh', test: /three-mesh-bvh/, priority: 60 },
                        // postprocessing must have higher priority than n8ao so n8ao's
                        // recursive deps (postprocessing library) don't get merged into the n8ao chunk
                        { name: 'postprocessing', test: /node_modules\/postprocessing\//, priority: 55 },
                        { name: 'postprocessing.ao', test: /node_modules\/n8ao\//, priority: 50 },
                        { name: 'rapier3d', test: /@dimforge\/rapier3d/, priority: 50 },
                        { name: 'materialx', test: /@needle-tools\/materialx/, priority: 50 },
                        { name: 'three-mesh-ui', test: /node_modules\/three-mesh-ui/, priority: 50 },
                        { name: 'three-quarks', test: /node_modules\/three\.quarks/, priority: 50 },
                        { name: 'peerjs', test: /node_modules\/peerjs/, priority: 50 },
                        { name: 'gltf-progressive', test: /\/gltf-progressive\//, priority: 40 },
                        // Independently skippable feature groups — only features large enough
                        // to justify a separate chunk (>50 KB). Smaller features merge into
                        // the main needle-engine chunk to reduce HTTP requests.
                        { name: 'needle-engine-ui', test: /engine-components\/ui\//, priority: 35 },
                        { name: 'needle-engine-webxr', test: /engine-components\/(webxr|avatar)\//, priority: 35 },
                        { name: 'needle-engine-particles', test: /engine-components\/particlesystem\//, priority: 35 },
                        // Postprocessing effects — separate so n8ao isn't eagerly loaded
                        { name: 'needle-engine.extras', test: /engine-components\/postprocessing\//, priority: 35 },
                        // Small 3rd-party deps used by the engine — merge into needle-engine
                        // to avoid tiny unnamed chunks.
                        { name: 'needle-engine', test: /node_modules\/websocket-ts\//, priority: 35 },
                        { name: 'needle-engine', test: /three-animation-pointer/, priority: 35 },
                        // register_types (codegen) must stay separate from lazy component code
                        // to preserve dynamic import() boundaries created by the treeshake plugin.
                        { name: 'register_types', test: /codegen\/register_types/, priority: 32 },
                        // Everything else (export, timeline, splines, web,
                        // physics, debug, utils, experimental) merges into the main
                        // needle-engine chunk.
                        { name: 'needle-engine', test: /@needle-tools\/engine/, priority: 30 },
                    ]
                };
            } else {
                rollupOutput.manualChunks = needleManualChunks;
            }
        }

        if (rollupOutput.chunkFileNames) {
            needleLog("needle-dependencies", "chunkFileNames already defined");
        }
        else {
            rollupOutput.chunkFileNames = handleChunkFileNames;
        }

        if (rollupOutput.assetFileNames) {
            needleLog("needle-dependencies", "assetFileNames already defined");
        }
        else {
            rollupOutput.assetFileNames = assetFileNames;
        }
    }


    // TODO: this was a test if it allows us to remove the sync import of postprocessing due to n8ao's import
    // config.build.rollupOptions.external = (source, importer, isResolved) => {
    //     if (importer?.includes("node_modules/n8ao/") || importer?.includes("node_modules/postprocessing/")) {
    //         console.log("EXTERNAL", importer);
    //         return true;
    //     }
    // }

    /** https://rollupjs.org/configuration-options/#output-assetfilenames
     * @param {import("vite").Rollup.PreRenderedAsset} chunkInfo */
    function assetFileNames(chunkInfo) {
        // "assets/..." is the default

        // this happens if e.g. a glTF or GLB file is link preloaded
        if (chunkInfo.name?.toLowerCase().includes(".glb") || chunkInfo.name?.toLowerCase().includes(".gltf")) {
            return `assets/[name][extname]`;
        }
        return `assets/[name].[hash][extname]`;
    }

    let needleEngineVersionedNameUsed = false;

    /** @param {import("vite").Rollup.PreRenderedChunk} chunk */
    function handleChunkFileNames(chunk) {
        if (chunk.name === 'needle-engine') {
            try {
                const version = tryGetNeedleEngineVersion();
                if (version && !needleEngineVersionedNameUsed) {
                    // Rolldown codeSplitting can create multiple chunks with the same name.
                    // Only assign the versioned name once — to the chunk with the most exports
                    // (the primary chunk). Additional same-named chunks get hash-based names.
                    const exportCount = chunk.exports?.length ?? 0;
                    if (exportCount > 10) {
                        needleEngineVersionedNameUsed = true;
                        const name = `assets/needle-engine@${version}.js`;
                        preloadScriptPaths.push(`./${name}`);
                        return name;
                    }
                }
            }
            catch (e) {
                needleLog("needle-dependencies", "Error reading version " + e, "warn");
            }
        }
        else if (chunk.name === 'three' || chunk.name === 'three-examples' || chunk.name === 'three-examples.extras') {
            // 'three' may be aliased to @needle-tools/three — check both package locations
            const version = tryGetPackageVersion("three") ?? tryGetPackageVersion("@needle-tools/three");
            const name = version
                ? `assets/${chunk.name}@${version}.js`
                : `assets/${chunk.name}.js`;
            // Preload three core and three-examples (loaders/libs) — both are always
            // needed. three-examples.extras is NOT preloaded (feature-specific).
            if (chunk.name !== 'three-examples.extras') {
                preloadScriptPaths.push(`./${name}`);
            }
            return name;
        }
        // else if(chunk.name === 'index') {
        //     console.log(chunk);
        //     debugger
        //     // this is the main chunk
        //     // we don't want to add a hash here to be able to easily import the main script
        //     return `index.js`;
        // }
        // Rename stray chunks that Rolldown extracted from dynamic import boundaries.
        // These get arbitrary names (e.g. "ScrollFollow", "src") but belong to engine deps.
        const ids = chunk.moduleIds ?? Object.keys(chunk.modules ?? {});
        if (ids.some(id => /engine-components/.test(id)) && !chunk.name.startsWith('needle-engine')) {
            return `assets/needle-engine.extras.[hash].js`;
        }
        // Rename generic "src" or "build" chunks — these are tiny re-export shims
        // Rolldown creates at dynamic import() boundaries. Give them a clearer name.
        if (chunk.name === 'src' || chunk.name === 'build') {
            return `assets/needle-engine.dep.[hash].js`;
        }
        return `assets/[name].[hash].js`;
    }

    /** 
     * @param {string} id 
     * @param {import('vite').Rollup.ManualChunkMeta | null} meta
    */
    function needleManualChunks(id, meta) {

        // console.log(id);

        if (id.includes("three/examples") || id.includes("three/addons")) {
            if (threeExamplesCoreRegex?.test(id)) {
                return "three-examples";
            }
            return "three-examples.extras";
        }
        if (id.includes('/three/')) {
            return "three";
        }

        if (id.includes("node_modules/n8ao/")) {
            detectSyncImports(id, meta);
            return "postprocessing.ao";
        }

        if (id.includes("node_modules/postprocessing/")) {
            // postprocessing bundle is preloaded https://github.com/vitejs/vite/pull/9938
            detectSyncImports(id, meta);
            return "postprocessing";
        }

        if (id.includes("@dimforge/rapier3d")) {
            detectSyncImports(id, meta);
            return "rapier3d";
        }

        if (id.includes("@needle-tools/materialx")) {
            detectSyncImports(id, meta);
            return "materialx";
        }

        if (id.includes("node_modules/three-mesh-ui")) {
            return "three-mesh-ui";
        }

        if (id.includes("node_modules/three.quarks")) {
            return "three-quarks";
        }
        // we want to bundle gltf-progressive separately because it also initiates Draco and KTX worker preloading
        if (id.includes("/gltf-progressive/")) {
            return "gltf-progressive";
        }
        if (id.includes("node_modules/needle-engine")
            // DEV
            // || id.includes("/package~/src/")
        ) {
            return "needle-engine";
        }
    }
}


/** 
 * @param {string} id 
 * @param {import('vite').Rollup.ManualChunkMeta | null} meta
*/
function detectSyncImports(id, meta) {
    if (meta) {
        if (meta.getModuleInfo) {
            const info = meta.getModuleInfo(id);
            if (info && info.importers.length > 0) {
                const isNeedleDev = id.includes("/package~/src/");
                if (isNeedleDev) {
                    console.warn(`WARN: SYNC IMPORTER DETECTED of ${id}`);
                    console.warn(info.importers)
                }
            }
        }
    }
}