mizan/backends/mizan-ts/src/ssr.ts

/**
 * SSR Bridge — manages a persistent Bun subprocess for React server-rendering.
 *
 * TypeScript port of `mizan-django/src/mizan/ssr/bridge.py`. Same wire
 * protocol: newline-delimited JSON-RPC over the worker's stdin/stdout, with
 * message-id correlation so concurrent renders don't cross.
 *
 *   → { "id": 1, "method": "render", "params": { "file": "/abs/Hello.tsx", "props": { ... } } }
 *   ← { "id": 1, "html": "<div>...</div>" }
 *   ← { "id": 1, "error": "..." }                                        (on failure)
 *
 * The worker (`workers/mizan-ssr/src/worker.tsx`) `import()`s the component file
 * and calls `renderToString` — no registry. It announces readiness with
 * `{ "id": 0, "ready": true }`; the bridge waits for that before accepting
 * renders, and restarts the worker if it exits.
 */

import { spawn, type ChildProcessWithoutNullStreams } from 'child_process'

export interface SSRBridgeOptions {
    /** Absolute path to the worker entry (workers/mizan-ssr/src/worker.tsx). */
    worker: string
    /** Per-render + startup timeout, seconds. Default 5. */
    timeout?: number
    /** Runtime to launch the worker. Default 'bun'. */
    runtime?: string
    /**
     * Args passed to the runtime before the worker path. Default `['run']`
     * (the Bun/`bun run <worker>` convention). Set `[]` for a runtime like
     * `node` that takes the script path directly.
     */
    runtimeArgs?: string[]
}

export interface RenderResult {
    html: string
}

interface Pending {
    resolve: (msg: any) => void
    reject: (err: Error) => void
    timer: ReturnType<typeof setTimeout>
}

export class SSRBridge {
    private readonly worker: string
    private readonly timeoutMs: number
    private readonly runtime: string
    private readonly runtimeArgs: string[]

    private proc: ChildProcessWithoutNullStreams | null = null
    private counter = 0
    private buffer = ''
    private readonly pending = new Map<number, Pending>()
    private readyPromise: Promise<void> | null = null
    private readyResolve: (() => void) | null = null
    private readyReject: ((err: Error) => void) | null = null

    constructor(options: SSRBridgeOptions) {
        this.worker = options.worker
        this.timeoutMs = (options.timeout ?? 5) * 1000
        this.runtime = options.runtime ?? 'bun'
        this.runtimeArgs = options.runtimeArgs ?? ['run']
    }

    private ensureRunning(): Promise<void> {
        if (this.proc !== null && this.proc.exitCode === null && this.readyPromise !== null) {
            return this.readyPromise
        }

        let settled = false
        this.readyPromise = new Promise<void>((resolve, reject) => {
            this.readyResolve = () => {
                if (!settled) {
                    settled = true
                    resolve()
                }
            }
            this.readyReject = (err) => {
                if (!settled) {
                    settled = true
                    reject(err)
                }
            }
        })

        const proc = spawn(this.runtime, [...this.runtimeArgs, this.worker], {
            stdio: ['pipe', 'pipe', 'pipe'],
        })
        this.proc = proc

        proc.stdout.setEncoding('utf-8')
        proc.stdout.on('data', (chunk: string) => this.onStdout(chunk))
        // Only react to THIS proc's exit — a stale exit event (from a worker we
        // already replaced) must not null out the freshly-spawned one.
        proc.on('exit', () => this.onExit(proc))
        proc.on('error', (err) => {
            this.readyReject?.(new Error(`SSR worker failed to spawn: ${err.message}`))
        })

        const startTimer = setTimeout(() => {
            this.readyReject?.(new Error(`SSR worker failed to start within ${this.timeoutMs}ms`))
            this.shutdown()
        }, this.timeoutMs)

        // Clear the start timer once ready settles (either way).
        this.readyPromise.then(
            () => clearTimeout(startTimer),
            () => clearTimeout(startTimer),
        )

        return this.readyPromise
    }

    private onStdout(chunk: string): void {
        this.buffer += chunk
        let nl: number
        while ((nl = this.buffer.indexOf('\n')) !== -1) {
            const line = this.buffer.slice(0, nl).trim()
            this.buffer = this.buffer.slice(nl + 1)
            if (!line) continue
            let msg: any
            try {
                msg = JSON.parse(line)
            } catch {
                continue // malformed line — ignore, matches the Python reader
            }
            this.onMessage(msg)
        }
    }

    private onMessage(msg: any): void {
        // Ready signal (id=0).
        if (msg.id === 0 && msg.ready) {
            this.readyResolve?.()
            return
        }
        const id = msg.id
        if (typeof id === 'number' && this.pending.has(id)) {
            const p = this.pending.get(id)!
            this.pending.delete(id)
            clearTimeout(p.timer)
            p.resolve(msg)
        }
    }

    private onExit(proc: ChildProcessWithoutNullStreams): void {
        // Ignore exit events from a worker we've already replaced.
        if (this.proc !== null && this.proc !== proc) return

        // Fail any in-flight requests; the next call re-spawns a fresh worker.
        const err = new Error('SSR worker exited')
        for (const [, p] of this.pending) {
            clearTimeout(p.timer)
            p.reject(err)
        }
        this.pending.clear()
        this.readyReject?.(err)
        this.proc = null
        this.readyPromise = null
    }

    private request(method: string, params: Record<string, any>): Promise<any> {
        const id = ++this.counter
        const frame = JSON.stringify({ id, method, params }) + '\n'

        return new Promise<any>((resolve, reject) => {
            const timer = setTimeout(() => {
                this.pending.delete(id)
                reject(new Error(`SSR ${method} timed out after ${this.timeoutMs}ms`))
            }, this.timeoutMs)
            this.pending.set(id, { resolve, reject, timer })

            try {
                this.proc!.stdin.write(frame)
            } catch (e: any) {
                this.pending.delete(id)
                clearTimeout(timer)
                reject(new Error(`SSR worker pipe broken: ${e?.message ?? e}`))
            }
        })
    }

    /** Render a React component file to HTML. Spawns the worker on first use. */
    async render(file: string, props: Record<string, any> = {}): Promise<RenderResult> {
        await this.ensureRunning()
        const msg = await this.request('render', { file, props })
        if (msg.error !== undefined) throw new Error(`SSR render failed: ${msg.error}`)
        return { html: msg.html }
    }

    /** Health check — resolves true when the worker answers a ping. */
    async ping(): Promise<boolean> {
        await this.ensureRunning()
        const msg = await this.request('ping', {})
        return msg.pong === true
    }

    /** Stop the Bun subprocess. */
    shutdown(): void {
        if (this.proc !== null) {
            try {
                this.proc.stdin.end()
            } catch {
                /* already closed */
            }
            try {
                this.proc.kill()
            } catch {
                /* already gone */
            }
            this.proc = null
            this.readyPromise = null
        }
    }
}