src/engine/wasm_engine.js

/* WASM backend — Swan inside the tab.
 *
 * Three pieces have to line up, and all three now come from one place:
 *
 *   1. DuckDB                 a real 1.5.4, built for Pyodide, from a wheel
 *   2. swan.duckdb_extension  the wasm_eh_pyodide build, supplying
 *                             `swan_compile_ast`
 *   3. pyrel_duckdb           the PyRel front-end — the DSL, the AST serializer
 *                             and the SQL post-processing are all Python
 *
 * All of it runs inside Pyodide, and there is exactly one database.
 *
 * They are not all loaded at once. ready() checks what is staged and
 * starts the worker; startPyrel() pulls Pyodide, numpy, pandas and the
 * front-end, roughly 100 MB and tens of seconds, so it is opt-in and reports
 * progress. Nothing here falls back to the sidecar on the caller's behalf:
 * which engine answered is what this backend exists to make visible.
 *
 * VERIFIED headlessly, against the same wasm binaries a browser runs:
 *   tools/verify_pyrel_wasm.mjs  a PyRel query compiles and returns rows, with
 *                                no CPython involved
 * And in real browsers: tools/smoke_browser.mjs and tools/smoke_templates.mjs
 * drive Chromium, tools/webkit_smoke.mjs drives WebKit.
 */

const ASSETS = {
  manifest: 'web/swan-wasm/manifest.json',
  python: 'web/swan-wasm/python/',
};

/* The compiled Swan extension, as the vendor script stages it. Named here
 * rather than read from the manifest because the probe below runs before the
 * manifest is trusted — checking a path the manifest supplied would make a
 * corrupt manifest look like a missing file. */
const PY_EXTENSION = 'swan.duckdb_extension.wasm';

/* The DuckDB the extension is built against. The manifest records what was
 * actually staged — read from the wheel's own filename — and probe() compares
 * the two, so a wheel swap is caught here rather than at LOAD swan. */
const REQUIRED_DUCKDB = '1.5.4';

async function assetExists(url) {
  try {
    return (await fetch(url, { method: 'HEAD' })).ok;
  } catch {
    return false;
  }
}

/** Hosts of a URL or an ordered chain of them, for saying where questions go. */
function hostOf(url) {
  const list = (Array.isArray(url) ? url : [url]).filter(Boolean);
  const hosts = list.map((u) => {
    try { return new URL(u).host; } catch { return String(u); }
  });
  // "a, then b" rather than a list: the order is the point, since the second is
  // only reached when the first refuses.
  return hosts.length > 1 ? `${hosts[0]}, then ${hosts.slice(1).join(', ')}` : hosts[0];
}

/* A short name for whichever generator answered, for the status line.
 *
 * "gpu" and "cpu" rather than the host, because that is the difference the user
 * actually feels -- one answers in seconds and the other in a minute or two --
 * and `maxdemarzi-black-swan-pyrel-gpu.hf.space` does not fit next to a row
 * count. Matched on the host so a chain pointed elsewhere still says something
 * true: an unrecognised host is reported as itself rather than guessed at.
 */
export function generatorLabel(host) {
  if (!host) return null;
  // The proxy is a *route to* the GPU, not a third kind of generator, so it
  // reports what the user gets rather than how it was reached — the hop is our
  // problem, not theirs.
  if (/-gpu\.hf\.space$/.test(host) || /-proxy\.hf\.space$/.test(host)) return 'gpu';
  if (/black-swan-pyrel\.hf\.space$/.test(host)) return 'cpu';
  if (/^(localhost|127\.0\.0\.1)(:|$)/.test(host)) return 'local';
  return host;
}

/* Name the generator that answered, in place, on the way back to the caller.
 *
 * The worker reports the host it reached; naming it is presentation, so it
 * happens here rather than in the worker. Left untouched when the worker said
 * nothing, so an older reply or a backend without a chain does not grow an
 * empty label. */
function labelVia(answer) {
  if (answer && answer.via && answer.via.host) {
    answer.via = { ...answer.via, label: generatorLabel(answer.via.host) };
  }
  return answer;
}

export function createWasmEngine({ onProgress = () => {}, nlq = null, nlqModel = null, nlqPrune = true } = {}) {
  let info = null;

  /* The PyRel front-end, started separately — see startPyrel(). */
  let worker = null;
  let pyrelReady = null;
  const inflight = new Map();
  let nextId = 1;

  /* Resolved by the worker's unsolicited `rules` message, tens of seconds after
   * the dataset is usable. Created up front so a caller that awaits rulesReady()
   * before the message arrives does not miss it. */
  let rulesResolve = null;
  const rulesPromise = new Promise((resolve) => { rulesResolve = resolve; });
  function settleRules(msg) {
    if (msg.error) { rulesResolve({ error: msg.error }); return; }
    if (info) info.rules = msg.rules;
    rulesResolve(msg.rules);
  }

  // A dead worker takes every outstanding call with it; failing them explicitly
  // beats leaving promises that never settle.
  function failAll(reason) {
    const err = new Error(`the Swan worker failed: ${reason}`);
    for (const [, p] of inflight) p.reject(err);
    inflight.clear();
  }

  /** One request/response round trip with the Worker. */
  function send(type, payload) {
    return new Promise((resolve, reject) => {
      const id = nextId++;
      inflight.set(id, { resolve, reject });
      worker.postMessage({ id, type, payload });
    });
  }

  async function probe() {
    onProgress('looking for the Swan wasm assets…');
    const res = await fetch(ASSETS.manifest).catch(() => null);
    if (!res || !res.ok) {
      throw new Error(
        'web/swan-wasm is not staged — run `node tools/vendor_swan_wasm.mjs`',
      );
    }
    const manifest = await res.json();
    if (manifest.duckdb !== REQUIRED_DUCKDB) {
      throw new Error(
        `the staged Swan extension targets duckdb ${manifest.duckdb}, but this engine expects ` +
        `${REQUIRED_DUCKDB} — re-run tools/vendor_swan_wasm.mjs`,
      );
    }
    /* Check the extension the worker actually loads.
     *
     * The worker writes python/swan.duckdb_extension.wasm into site-packages,
     * where swan finds it by scanning sys.path. Checking anything else means
     * guarding a file no code reads while the one every page depends on goes
     * unchecked, which is exactly what this used to do. */
    if (!(await assetExists(`${ASSETS.python}${PY_EXTENSION}`))) {
      throw new Error(
        'the Swan extension is not staged (web/swan-wasm/python/) — re-run tools/vendor_swan_wasm.mjs',
      );
    }
    return manifest;
  }

  return {
    name: 'wasm',

    async ready() {
      const manifest = await probe();

      /* The worker owns the database. Everything — sql(), schema(), pyrel() —
       * goes through it, so there is exactly one DuckDB and one Swan catalog.
       *
       * Two databases would mean data loaded through one being invisible to
       * the other, with nothing to say so. The Python side needs a synchronous
       * API, which cannot be driven from this thread anyway. */
      /* A SharedWorker, so the engine outlives the page that started it.
       *
       * Every page here is a separate document -- this is a static site, not an
       * SPA -- so a dedicated Worker died with whichever page created it, taking
       * Pyodide, DuckDB, the model and all 117 evaluated rules with it.
       *
       * Measured on Chromium, what this does and does not buy:
       *
       *   a second tab, first still open   31 s -> 0.3 s, and it answers
       *   navigate away and come back      no better -- the worker is gone
       *
       * The second row is a property of SharedWorker, not a bug here: one lives
       * only while at least one document is connected, and most of this site
       * never connects -- Galaxy, Ontology and Ladder read a static graph.json
       * and have no engine at all. Leaving Horizon drops the owner count to zero
       * and the browser collects it.
       *
       * A keepalive port on those pages was tried and did not save it; the
       * handoff still has a moment with no owner. It was removed rather than
       * shipped, because an optimisation that did not work under measurement is
       * worse than none -- it reads as covered.
       *
       * What would actually survive navigation is persisting the DuckDB file to
       * IndexedDB and restoring on boot, which would survive closing the browser
       * too. Pyodide's own ~10 s boot would remain.
       *
       * `worker` is the port rather than the worker: postMessage/onmessage look
       * the same on it, so everything downstream is unchanged. */
      const shared = new SharedWorker(new URL('./pyrel-worker.js', import.meta.url), { type: 'module' });
      worker = shared.port;
      shared.onerror = (e) => failAll(e.message || 'unknown error');
      worker.onmessage = (e) => {
        const msg = e.data || {};
        if (msg.type === 'progress') { onProgress(msg.text); return; }
        /* The rules tier finishes long after loadDataset() resolves, so it
         * arrives unsolicited rather than as a reply. Kept as a promise so a
         * caller can await it without knowing that. */
        if (msg.type === 'rules') { settleRules(msg); return; }
        const pending = inflight.get(msg.id);
        if (!pending) return;
        inflight.delete(msg.id);
        if (msg.ok) { pending.resolve(msg.result); return; }
        /* A failed question still carries the program that failed and the
         * translated database error. Carrying them on the Error is what lets the
         * page show the PyRel it tried next to why it did not run — the
         * difference between "it didn't work" and a query you can fix by hand in
         * the box below. horizon.js reads `err.code`. */
        const failure = new Error(msg.error);
        if (msg.code) failure.code = msg.code;
        if (msg.detail) failure.detail = msg.detail;
        pending.reject(failure);
      };
      /* A port does not surface worker errors -- only the SharedWorker object
       * does, via onerror above -- but a messageerror still lands here, and a
       * failure that reaches neither would leave every caller hanging. */
      worker.onmessageerror = () => failAll('the worker sent a message this page could not read');
      worker.start();

      onProgress('starting the Swan worker…');
      /* `init` answers with whatever this shared engine has already done, which
       * on a second page is usually everything. */
      const boot = await send('init');
      if (boot.rules) settleRules(boot.rules.error ? { error: boot.rules.error } : { rules: boot.rules });

      /* Adopt the shared worker's state as this page's own.
       *
       * `pyrelReady` gates pyrel(), and it is per-page: without this, a page that
       * returned to an engine another page had already started would show the
       * loaded model, accept a query, and then refuse it with "the PyRel
       * front-end is not started" -- the page's local flag disagreeing with the
       * worker it is attached to. */
      if (boot.pyrelStarted) pyrelReady = Promise.resolve(boot.python || null);

      info = {
        swan: manifest.swan,
        duckdb: manifest.duckdb,
        /* Not null when another page already loaded it. horizon.js keys the
         * "Load it" invitation off this, so a returning page goes straight to
         * a queryable model instead of being offered one it already has. */
        model: boot.dataset?.model || null,
        rules: boot.rules || null,
        /* True once the page has been told where to send a prompt. The pipeline
         * itself is already here — see ask() — so this reports configuration
         * rather than capability. Without an endpoint it stays false and `auto`
         * behaves exactly as it did, falling through to the sidecar for a page
         * that needs natural language. */
        nl: !!nlq,
        /* Available but not loaded: startPyrel() brings up the Python front-end
         * on demand, because it is ~100 MB. `auto` reads this flag, so leaving it
         * false is what makes a page needing PyRel fall through to the sidecar
         * rather than settling on an engine that would refuse on first use. */
        pyrel: !!boot.pyrelStarted,
        pyrelAvailable: !!manifest.python,
        /* Say where questions go, because by default they go somewhere. Every
         * other claim this page makes is "nothing leaves the tab", so the one
         * thing that does has to be visible rather than discovered. */
        note: manifest.python
          ? (nlq
            ? `Swan runs in the browser; questions go to ${hostOf(nlq)}`
            : 'Swan runs in the browser; PyRel available via startPyrel()')
          : 'Swan runs in the browser; the Python front-end is not staged',
      };
      return info;
    },

    /** Raw SQL, executed in the worker with Swan's optimizer loaded. */
    async sql(text) {
      const { columns, rows } = await send('sql', text);
      return { code: text, rows: rows.map((r) => Object.fromEntries(columns.map((c, i) => [c, r[i]]))) };
    },

    schema() { return send('schema'); },

    /**
     * Add the PyRel front-end — Pyodide, numpy, pandas and pyrel_duckdb — on top
     * of the database ready() already opened in the worker.
     *
     * Separate from ready() and never automatic, because it is roughly 100 MB
     * and tens of seconds. A page that paid that on open would look broken
     * rather than slow, so the caller decides when — and gets `onProgress`
     * throughout, since silence for that long is indistinguishable from a hang.
     *
     * It lands in the same worker and on the same connection, so a model built
     * here sees data loaded through sql() and vice versa.
     */
    async startPyrel({ onProgress: report = onProgress } = {}) {
      if (pyrelReady) return pyrelReady;
      if (!worker) throw new Error('engine not ready — await ready() first');
      if (!info?.pyrelAvailable) {
        throw new Error(
          'the Python front-end is not staged — re-run tools/vendor_swan_wasm.mjs with ' +
          'pyodide installed and a .venv containing rai-swan',
        );
      }
      pyrelReady = (async () => {
        report('starting the PyRel front-end (~100 MB, first time only)…');
        const result = await send('python');
        info.pyrel = true;
        /* Say where questions go, because by default they go somewhere.
         *
         * Everything else about this engine is "nothing leaves the tab", which
         * is the claim the page is built on -- so the one thing that does leave
         * has to be visible rather than discovered. The host, not the full URL:
         * enough to know where it went, short enough to sit in a header. */
        info.note = nlq
          ? `Swan and PyRel run in the browser; questions go to ${hostOf(nlq)}`
          : 'Swan and PyRel both run in the browser';
        return result;
      })();
      return pyrelReady;
    },

    /** Define the model inside the Worker, from PyRel source. */
    loadModel(source) { return send('model', source); },

    /**
     * Load the staged demo dataset and the model describing it.
     *
     * Without this the engine has Swan and an empty database — a demo that
     * cannot demo. Separate from startPyrel() because the two costs are very
     * different: the runtime is ~100 MB, the data is ~2 MB, and a caller may
     * reasonably want one without the other.
     */
    async loadDataset() {
      await this.startPyrel();
      const result = await send('dataset');
      // Reflect it on the info object callers already hold, so a UI that
      // rendered "no model" at boot can show the concept count without
      // re-reading anything.
      if (info) info.model = result.model;
      return result;
    },

    /** Resolves when the rules tier finishes, which is well after loadDataset().
     *  `{error}` if it could not be evaluated at all. */
    rulesReady() { return rulesPromise; },

    /**
     * Re-read the model and recompute every rule against the data as it stands.
     *
     * A warm load declares the rules and reuses values a previous session
     * computed. This is the way back to computing them, for when those values
     * are not trusted. Slow on purpose — it is the full rules run.
     *
     * Note that rulesReady() has long since settled by the time anyone can press
     * this, so it does not re-arm; the fresh tally comes back in the reply, and
     * `info.rules` is updated for anything reading it later.
     */
    async reloadModel() {
      const result = await send('reloadModel');
      if (info) { info.model = result.model; info.rules = result.rules; }
      return result;
    },

    /**
     * Forget the cached database, then stand the worker down.
     *
     * The caller reloads the page afterwards — that is not optional, it is how
     * the data gets rebuilt. Everything this worker holds goes with it, which is
     * the point: a rebuild in place would have to pull the database out from
     * under a live model.
     */
    reloadData() { return send('reloadData'); },

    async pyrel(code) {
      if (!pyrelReady) {
        /* Not proxied to the sidecar on the caller's behalf. Quietly forwarding
         * would report "wasm" for a query that ran elsewhere, and which engine
         * answered is the one thing this engine exists to make visible. */
        throw new Error(
          'the PyRel front-end is not started — call startPyrel() first (it is a ~100 MB download), ' +
          'or use the server engine',
        );
      }
      await pyrelReady;
      return send('pyrel', code);
    },

    /**
     * Ask in English, with generation done by whatever `?nlq=` points at.
     *
     * The pipeline is entirely in the browser — schema, few-shot selection,
     * sandbox, error translation, retries — because `pyrel_duckdb.nlq` ships in
     * the compiled module. Only the text has to come from somewhere, and
     * `answer_question` takes that as an injected function. So this engine is no
     * longer barred from natural language by not holding credentials; it is
     * barred only by not being told where to send a prompt.
     */
    async ask(question) {
      if (!nlq) {
        throw new Error(
          'the wasm engine has no generation endpoint — load the page with '
          + '?nlq=<url> to point it at a model, or use the server engine',
        );
      }
      await this.startPyrel();
      return labelVia(await send('ask', { question, endpoint: nlq, model: nlqModel, prune: nlqPrune }));
    },

    /* ---- templates ------------------------------------------------------
     *
     * Swan's templates are notebooks: prose, then a cell, then the next cell
     * using what the last one built. So this is two calls rather than one —
     * prepare the data and a namespace, then run cells into it — and the page
     * decides when each cell runs.
     *
     * Deliberately not part of ask()/pyrel(). Those answer questions about the
     * *loaded* model; a template brings its own model, its own data and its own
     * solver, and shares nothing with the TPC-DS engine but the Pyodide runtime.
     */
    async prepareTemplate({ slug, database, tables, packages }) {
      // startPyrel(), not loadDataset(): a template needs the Python front-end
      // but none of TPC-DS, and making a visitor wait ~100 MB for data no
      // template reads would be the wrong trade entirely.
      await this.startPyrel();
      return send('templatePrepare', { slug, database, tables, packages });
    },

    /** One notebook cell, into the namespace prepareTemplate() opened. */
    async runTemplateCell(source) {
      return send('templateCell', { source });
    },

    /** The model those cells built, as a three-tier graph the ladder can read.
     *  `{graph, models}` — `graph` is null when the cells built no Model at
     *  all, which is the honest answer for a template that only queries. */
    async templateGraph() {
      return send('templateGraph', {});
    },

    /** Tell the worker this page is done with its session.
     *
     *  The worker cannot detect a closed page — a SharedWorker port gives no
     *  reliable signal, and inferring one from a failed postMessage tore down
     *  a live tab's namespace. So the page says so itself, and the worker
     *  bounds the rest. */
    async closeTemplate() {
      return send('templateClose', {});
    },

    /** ask() with the conversation so far — `[{question, code}]`, held by the
     *  caller for the same reason the sidecar does not hold it: two tabs must
     *  not capture each other's context. */
    async chat(question, history = []) {
      if (!nlq) {
        throw new Error(
          'the wasm engine has no generation endpoint — load the page with ?nlq=<url>',
        );
      }
      await this.startPyrel();
      return labelVia(await send('ask', { question, history, endpoint: nlq, model: nlqModel, prune: nlqPrune }));
    },
  };
}