Vite Worker Plugin

vite-client-ext.d.ts

declare module '*?wrkr' {
    export const url: string;
    export const create: {
        (): Worker;
    };
}
declare module '*?srvcwrkr' {
    export const url: string;
    export const register: {
        (
            options?: Omit<RegistrationOptions, 'type'>,
        ): Promise<ServiceWorkerRegistration>;
    };
}
              

worker.ts

import debug from 'debug';
import type {
    LoadResult,
    ResolvedId,
    ResolveIdResult,
    TransformResult,
} from 'rollup';
import * as rollup from 'rollup';
import type {Plugin, ResolvedConfig} from 'vite';

const spacesRe = /^(?<spaces> +)[^ ]*/u;
function trimCode(code: string): string {
    let spacesReplaceRe: RegExp | undefined = undefined;
    const lines = code.split('\n');
    for (const line of lines) {
        const result = line.match(spacesRe);
        if (result?.groups === undefined) {
            continue;
        }
        spacesReplaceRe = new RegExp(
            `^[ ]{${result.groups.spaces.length}}`,
            'u',
        );
    }
    if (spacesReplaceRe === undefined) {
        return code;
    }
    // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
    return lines.map((line) => line.replace(spacesReplaceRe!, '')).join('\n');
}

type PluginForwardFilter = (plugins: Plugin[]) => Plugin[];
type WorkerGenerator = (url: string, type: 'classic' | 'module') => string;

function generateWorkerCreator(
    url: string,
    type: 'classic' | 'module',
): string {
    return trimCode(`
        export const url = ${url};
        export function create() {
            return new Worker(${url}, {type: "${type}"});
        }`);
}

function generateServiceWorkerCreator(
    url: string,
    type: 'classic' | 'module',
): string {
    return trimCode(`
        export const url = ${url};
        export function register(options) {
            return navigator.serviceWorker.register(${url}, {...options, ...{type: "${type}"}})
        }`);
}

export function workerPlugin(filter?: PluginForwardFilter): Plugin {
    const log = debug('vite-plugin-wrkr');
    let config: ResolvedConfig;
    let isDev = true;
    const ids = new Map<string, WorkerGenerator>();
    const urlQueryRe = /(?<path>.*)\?(?<type>wrkr|srvcwrkr)$/u;

    return {
        name: 'vite-plugin-wrkr',
        enforce: 'pre',

        configResolved(config_): void {
            config = config_;
            isDev = config_.command === 'serve';
        },

        async resolveId(source, importer, options): Promise<ResolveIdResult> {
            if (isDev) {
                // Look at the URL parameters and check if we need to handle it
                const result = source.match(urlQueryRe);
                if (result?.groups === undefined) {
                    return null;
                }
                let generator: WorkerGenerator;
                switch (result.groups.type) {
                    case 'wrkr':
                        generator = generateWorkerCreator;
                        break;
                    case 'srvcwrkr':
                        generator = generateServiceWorkerCreator;
                        break;
                    default:
                        throw new Error(
                            `Unknown wrkr type ${result.groups.type}`,
                        );
                }

                // eslint-disable-next-line @typescript-eslint/ban-types
                let resolved: ResolvedId | null = null;
                // Resolve source path first
                resolved = await this.resolve(result.groups.path, importer, {
                    ...options,
                    skipSelf: true,
                });
                if (resolved === null) {
                    return null;
                }

                // Dispatch resolved to asset plugin via <resolved>?url
                resolved = await this.resolve(`${resolved.id}?url`, importer, {
                    ...options,
                    skipSelf: true,
                });
                log('resolved', source, importer, '->', resolved);
                if (resolved === null) {
                    return null;
                }
                ids.set(resolved.id, generator);
                return resolved;
            } else {
                return null;
            }
        },

        load(id): LoadResult {
            if (isDev) {
                if (!ids.has(id)) {
                    return null;
                }
                log('load', id);
                return null;
            } else {
                // Look at the ID URL parameters and check if we need to handle it
                const result = id.match(urlQueryRe);
                if (result?.groups === undefined) {
                    return null;
                }
                let generator: WorkerGenerator;
                switch (result.groups.type) {
                    case 'wrkr':
                        generator = generateWorkerCreator;
                        break;
                    case 'srvcwrkr':
                        generator = generateServiceWorkerCreator;
                        break;
                    default:
                        throw new Error(
                            `Unknown wrkr type ${result.groups.type}`,
                        );
                }
                ids.set(id, generator);
                log('load', id);
                return '';
            }
        },

        async transform(code, id): Promise<TransformResult> {
            const generator = ids.get(id);
            if (generator === undefined) {
                return null;
            }
            log('transform', id, code);

            if (isDev) {
                // Retrieve URL from generated code of the asset plugin
                // TODO: Ugly workaround as we can't seem to import `fileToUrl`
                //       from 'vite/src/node/plugins/asset'
                const url = code.replace('export default ', '');
                return {
                    code: generator(url, 'module'),
                };
            } else {
                // Bundle the file as a new entry
                const bundle = await rollup.rollup({
                    input: id,
                    plugins: (filter?.([...config.plugins]) ??
                        config.plugins) as rollup.Plugin[],
                });
                try {
                    const {
                        output: [chunk],
                    } = await bundle.generate({
                        format: 'iife',
                        sourcemap: config.build.sourcemap,
                    });
                    const url = `"__VITE_ASSET__${this.emitFile({
                        type: 'asset',
                        name: chunk.fileName,
                        source: chunk.code,
                    })}__"`;
                    return {
                        code: generator(url, 'classic'),
                    };
                } finally {
                    await bundle.close();
                }
            }
        },
    };
}