Compile models from Stan source as a preload step (#137)

[kiante-fernandez]

Follow-up to the controller API (#141, now merged). Rebased onto main. The first commit is the preload feature (#137); the rest are measured, readability-oriented cleanup — dead code, facade slim, lazy debug UI, tighter interfaces, comment retouch, and a README/docs pass (details at the bottom). Supersedes the draft #144 (opened stacked on #141 before it merged).

Summary

Implements #137: a model may be supplied as Stan source (stanCode) instead of committed moduleUrl/wasmUrl artifacts, and is compiled to WASM during an explicit preload step — the same pattern as jsPsychPreload gates media, applied to model compilation.

const stanCode = await fetch("./exponential.stan").then((r) => r.text());

const ado = jsPsychADO.createController(jsPsych, {
  model: { id, stanCode, params, designKeys, responseSpace, stanData, responseProb },
  design_grid,
  // compile: { server } — defaults to the public stan-playground compile server
});

jsPsych.run([
  instructions,
  ado.preload(),                 // "Preparing the experiment…" gate; usually instant (see below)
  ...ado.createTimeline(trial),
  end,
]);

No offline compile step, no committed artifacts, no patch:wasm, and no hand-written prior (derived from the Stan source via parseStanPriors). Closes #137.

Design

• Eager compilation, handle-owned worker. The compile request fires at createController, so it overlaps the welcome/instruction screens. The readiness chain is: POST source → content-addressed model_id → download + verify the compiled glue → the handle's shared Stan worker imports the module and instantiates its wasm. ado.ready() / ado.preload() await that whole chain — so a green preload certifies the model is actually loadable, not merely compiled (see the review-response note below; the browser-smoke preload wait is ~54 ms, dominated by the real worker load).
• ado.preload(opts) / ado.ready(). The gate trial is a ~40-line self-contained plugin (no new plugin dependency). On failure it renders the compiler's actual error message — a student's stanc syntax error is meant to be read — and aborts visibly; max_load_time gives it a jsPsychPreload-style deadline. The trial is optional: without it the first posterior update simply awaits readiness. ado.ready() is public for custom loading UI.
• Mock-mode handles stay network-free. controller: "mock" never contacts the compile server (an offline dev loop can't be killed by a service it doesn't need); a per-timeline controller: "stan" override starts the compile lazily.
• Validation sharpens the package contract: moduleUrl XOR stanCode; a leftover wasmUrl on a source model is rejected (it would pair the server-compiled glue with a stale local binary); non-object priors on source models are rejected.
• Caching economics. The compile server's model_id is a content hash: the first request for a given source ever pays the compile; every identical request — i.e. all participants — is a cache hit plus a ~1 MB artifact download.

Deployment reality (documented, not hidden)

This is the dev/teaching/prototyping path. The public compile server currently allows browser requests only from its own allowlisted origins (works today from 127.0.0.1:3000); deployed experiments need a self-hosted compile server (docker run -p 8083:8080 ghcr.io/flatironinstitute/stan-wasm-server:latest) or an upstream CORS change (see #137 for the analysis). Committed artifacts remain the production path — compile once, commit main.js + main.wasm — which is also the right answer for reproducibility. Graduating between the two is swapping stanCode for moduleUrl/wasmUrl in the same model object.

Demo

demos/byo_model_exponential/from_source.html — the bring-your-own-model demo with zero offline compile steps, sharing responseProb/stanData with the committed-artifacts variant so the two pages can never fit different models.

Validation

• Unit tests covering eager compile, prior derivation, custom server, the worker receiving the compiled URL, preload finish/error/timeout paths, mock network-free, artifact-download failure, session cache, and wasmUrl/prior validation (plus the review-response coverage below).
• A browser smoke driving the full chain — compile POST → preload gate → artifact download → worker init → real WASM inference — against a local mock compile server (CI never depends on the external service), wired into test:browser.
• The end-to-end path against the real public compile server was verified manually from an allowlisted origin (compile cache hit, correct application/wasm MIME, posterior recovery).
• Types (ModelPackage.stanCode, CompileConfig, preload/ready on the handle) + CHANGELOG + README.

Review response (@githubpsyche)

Three review comments, all addressed — with an explanatory comment on each in the thread:

1. prepareModel(...) artifact pairing. Fixed at the root: the source-shape + no-wasmUrl-on-source rule is unified behind one stanUrl-aware validateSourceSpec shared by validateModel and prepareModel (one canonical message, not three drifting copies). This also closed two latent stanUrl gaps in the public validateModel, and createController now rejects a stanUrl-only model with an actionable "compile with prepareModel first" message.
2. ado.preload() didn't cover the worker load. Deep fix: Stan worker ownership moved from the per-timeline controller to the handle (one lazy, shared worker). ado.ready() / preload() now await the worker's module import + wasm instantiation — for committed models too — so a green preload certifies loadability (browser-smoke preload wait went ~2 ms → ~54 ms). createController stays worker-free; practice→main timelines reuse one worker (inited once); a compile/download/worker-load failure rejects ready()/preload and the first update.
3. "in-browser compile" wording. Reworded the docs and demo text to "compile server" — compilation is server-backed.

The #2 design was chosen over eager-init and source-only alternatives (handle-owned lazy-shared worker is the only one that certifies both committed and source models while keeping createController side-effect-light); #1 and #2 were both adversarially reviewed. Unit suite + 10 browser smokes + bundler + recovery/parity smokes green.

Also in this PR: measured codebase cleanup

The second commit (d584fec) is a round of readability-oriented cleanup bundled alongside the feature — no behavioral or API change beyond one documented model-export trim, so it can be reviewed quickly and separately from the preload work:

• Dead code removed — the unreferenced compile_stan_model.js, fully superseded by prepareModel.
• Facade slimmed (index.js ~780 → ~695 lines) — outcome-label resolution and the debug-flag/?debug resolver extracted to small focused modules (response_labels.js, debug_flag.js).
• Debug UI lazy-loaded — the ~1,400 lines of chart/SVG/panel code under src/ado/debug/ now load only when ?debug is set, so a production bundler splits them into a chunk participants never download (~18 KB off the entry bundle). Debug-on behavior is unchanged.
• Tighter module interfaces — dropped over-exports that were only used internally; a model's public interface is cleanly its default export (the duplicate/aliased named exports are gone; standalone math helpers stay). This is the one public-surface change, noted in the CHANGELOG.
• Comment retouch, repo-wide — rewrote provenance/history comments (bare issue tags, "fixed"/"used to…"/"the old…" phrasings) to describe present behavior instead of how the code changed, keeping only a few genuine pointers to external context (compile-server CORS Compile models from Stan source as a preload step (jsPsychPreload-style) #137, the bundler wasm-resolution path Make the committed WASM survive production bundlers (Vite/webpack) for npm distribution #57, upstream Vite #10837). Comment/doc-string/test-name strings only — no executable code changed.
• README + docs pass — cut the README to a ~200-word landing page (overview, a minimal end-to-end API example, doc links) and moved the full reference (API, bundler setup, adaptive stopping, internals) into docs/usage.md. No content dropped; all links, demo paths, and npm scripts verified.

Fully re-validated after the cleanup: 206 unit tests, typecheck, prettier, 8 real-WASM recovery + parity smokes, the bundler smoke (chunk split intact), and all 10 browser smokes (every demo, debug on and off).

[githubpsyche]

Next I think there may be one artifact-pairing issue in the lower-level prepareModel(...) path.

Reviewing my understanding...

There are two ways a Stan model can reach createController(...).

The normal deployed path is a precompiled model package. In that case the model object points at committed artifacts:

moduleUrl: new URL("./main.js", import.meta.url).href,
wasmUrl: new URL("./main.wasm", import.meta.url).href,

Those two files are a matched pair: the worker imports main.js, and that generated JS loads the matching main.wasm.

The source/prototyping path is different. A user supplies Stan source code (stanCode or stanUrl), and prepareModel(...) sends that source to a compile server. The compile server returns a new compiled main.js, whose matching .wasm lives next to it on the compile server. In that path, any .wasm URL that was already present on the input object should not survive, because it may refer to an older local compile or a different model.

It looks like the direct createController(...) source-code path handles this carefully: after prepareModel(...) resolves, it passes the server-returned moduleUrl forward with wasmUrl: null.

Now the possible gap...

The exported lower-level prepareModel(...) helper strips stanCode / stanUrl, but returns the rest of the input object unchanged:

const { stanCode: _code, stanUrl: _url, ...rest } = spec;
return { ...rest, prior, moduleUrl };

If rest includes a stale wasmUrl, the prepared model can become:

{
  moduleUrl: "https://compile-server/download/new-model/main.js",
  wasmUrl: "https://some-old-or-local/main.wasm"
}

Then createController(...) sees a model with moduleUrl and no stanCode, so it treats it like a normal precompiled model and forwards both URLs to the worker. That can pair server-compiled JS glue with an unrelated wasm file.

I think prepareModel(...) should enforce the same invariant as the main source-code path: either reject wasmUrl when the input is stanCode / stanUrl, or strip it before returning the prepared model. A small regression test on prepareModel(...) itself would catch this, since the direct createController(...) path already avoids the mismatch.

[githubpsyche]

Next I think there may be one readiness issue around ado.preload().

Reviewing my understanding from scratch...

When a user supplies Stan source code instead of committed main.js / main.wasm files, the browser sends that Stan source to a compile server. The compile server returns a generated main.js. Later, the Stan worker imports that generated main.js, and the generated JS loads the matching main.wasm.

As I understand it, ado.preload() is the optional jsPsych trial a study author can place before the adaptive timeline. It gives participants a waiting screen while the source model is prepared, then records preload data and lets the experiment continue.

That participant-facing role seems important. If the preload trial succeeds, I would expect the model-loading path needed for adaptive trials to have been checked.

Now the possible gap...

Right now, ado.preload() appears to wait for compilation and download of the generated main.js, but not for the Stan worker to actually load the compiled model. In src/index.js, the readiness chain compiles the model and fetches/consumes the returned main.js. But worker initialization still happens later, when the Stan controller starts:

model_ready = module_source.then(({ moduleUrl, wasmUrl }) => client.init(moduleUrl, wasmUrl));

That means the preload trial can finish with ado_preload_ok: true before the worker import / wasm-load path has actually been exercised.

I think preload should include that worker initialization step. The code comments already describe ready() / preload as promising that the model is “USABLE,” and the preload row records ado_preload_ok. To me, that success flag should mean the participant can proceed into adaptive trials without immediately discovering that the compiled model was not actually loadable by the worker.

I realize this may not be a one-line preload-plugin change, because worker initialization currently happens inside the Stan controller created for the adaptive timeline. But I think the readiness promise should eventually cover that same initialization step, whether by moving the initialization earlier, sharing the initialization promise with preload, or otherwise making the preload success flag correspond to the worker-loaded model state. If that is too invasive for this PR, then narrowing the comments/data wording would at least make the behavior explicit, though that seems like the less useful outcome.

[githubpsyche]

Reviewing my understanding...

There are two ways a user can provide a Stan model.

The deployed/self-contained path is to precompile the model and ship the generated main.js / main.wasm files with the experiment. In that case the experiment can run from static assets, without contacting a compile server during the participant session.

The Stan-source path is different. A user supplies stanCode or stanUrl; the browser sends that Stan source to a compile server; the server returns a generated main.js URL; and the browser/worker then loads the returned compiled artifacts. So the browser initiates the preparation, but the Stan-to-WASM compilation itself is server-backed.

Now the possible gap...

A few docs/demo strings describe the source path as “in-browser compile” or “compiles it in the browser.” For example:

• docs/usage.md: ado.preload(...) waits for the “in-browser compile.”
• demos/byo_model_exponential/README.md: from_source.html “compiles it in the browser.”
• demos/byo_model_exponential/from_source.html: “This demo compiles its Stan model in the browser while you read this.”

I think that wording could confuse users about the deployment requirements. For source-supplied models, they need access to an allowlisted public compile server or a self-hosted compile server. The no-compile-server path is the committed-artifacts path.

The surrounding docs already mention compile servers and CORS, so I think this can probably be fixed with wording rather than a larger docs change. Maybe use “compiled through a compile server” in docs, and “Preparing the model...” in participant-facing demo text?

[kiante-fernandez]

@githubpsyche — re: the prepareModel(...) artifact-pairing gap. Addressed at the root in 5d98127.

Rather than only guarding prepareModel, I pulled the source-spec contract that prepareModel and validateModel had each been hand-rolling into one shared, stanUrl-aware validator:

validateSourceSpec(spec) in validation.js — pure/synchronous, returns problem strings. It enforces "exactly one of moduleUrl | stanUrl | stanCode" and the "no wasmUrl on a source spec" rule via a single canonical message. validateModel pushes its problems through err(); prepareModel throws the first. So the rule and its wording now live in one place instead of three copies that had already drifted apart.

Unifying the two validators surfaced two latent stanUrl blind spots in validateModel, both now fixed:

• Shape: a { stanUrl } spec was misread as "neither moduleUrl nor stanCode" — validateModel never considered stanUrl.
• wasmUrl: { stanUrl, wasmUrl } slipped past the wasmUrl rule, because its gate keyed on the stanCode-only isSourceModel. prepareModel caught it; validateModel didn't. Both now go through the same seam.
• I also made the prior-omission exemption stanUrl-aware (prepareModel derives the prior from a fetched stanUrl exactly as from inline stanCode), so a prior-less stanUrl spec is a valid package.

isSourceModel deliberately stays stanCode-only: createController derives the prior synchronously and can't fetch a stanUrl, so treating stanUrl as inline source there would crash prior derivation. Instead createController now rejects a stanUrl-only model with an actionable "compile with prepareModel first" message, rather than surfacing later as a Stan-init crash.

New tests: a validateSourceSpec shape/wasmUrl matrix, stanUrl regression pins for both blind spots, and the createController guard. Full unit suite + compile-preload smoke green.

(Your other two points are handled too: the ado.preload() / worker-init readiness gap — I narrowed the wording so ado_preload_ok honestly means compiled + downloaded, with the full worker-init-in-preload as a follow-up; and the "in-browser compile" wording — reworded to "compile server" in the docs and demo. Happy to split those into their own notes if useful.)

[kiante-fernandez]

@githubpsyche — circling back on your ado.preload() readiness point. Earlier I only narrowed the wording so ado_preload_ok honestly meant "compiled + downloaded"; I've now done the deeper fix so preload actually covers the worker load (301b368, tidied in b05a184).

The change: worker ownership moved from the per-timeline controller to the handle. Previously each createStanAdoController created its own worker client and called client.init() inside start() — which runs after the preload gate — so ado_preload_ok could be true before the worker had imported the module or instantiated wasm. Now:

• The handle owns one shared worker via a lazy, memoized ensureStanRuntime(); ado.ready() / ado.preload() await its client.init(). So a green preload means the worker has genuinely imported the compiled module and instantiated its wasm — for committed models too, not just compiled ones (that was the deciding factor: committed wasm is the production default, so a source-only fix would have left the common case uncertified). The browser smoke shows it working — the preload wait went from ~2 ms (compile+download only) to ~54 ms (now including the real worker load).
• The controller is now a thin adopter: it receives the shared worker_client + worker_ready and just awaits them; start() no longer owns any worker lifecycle.
• createController stays worker-free until readiness is awaited (so validation-only use and the ?debug dev loop don't spin up a worker), and practice→main timelines reuse the same worker (inited once — there's a test asserting exactly one init across both).

Failure propagation is preserved on both paths: a compile / download / worker-load failure rejects ready() (preload renders it + aborts with ado_preload_ok: false) and the first update() (mid-run abort). New tests cover the worker-load-failure path and the committed-model init (asserting the bundler wasmUrl, #57, reaches the worker).

One honest caveat: ready() reflects the handle-level controller. A controller: "mock" handle resolves immediately (no wasm), and a per-timeline controller: "stan" override on a mock-default handle is not gated by preload() — that timeline's worker loads when it's built. It's documented in ready()'s doc comment; put such an override on a stan-default handle if you want preload to cover it.

I picked the "handle owns a lazy shared worker" shape over eager-init and source-only alternatives (it's the only one that certifies both committed and source models while keeping createController side-effect-light). Validated across the unit suite, all 10 browser smokes, the bundler smoke, and recovery/parity smokes.

[githubpsyche]

I think there may be a validation mismatch in the documented prepareModel(...) → createController(...) workflow. Before suggesting a resolution, I want to check that I understand how the two supported artifact-loading paths are intended to work.

Reviewing my understanding...

Before a Stan model can run in the browser, it must be compiled into two matching files: main.js, which loads the model, and main.wasm, which contains the compiled model itself.

For the production path, those files are shipped with the experiment. Build tools such as Vite and webpack may rename main.wasm, so the model includes an explicit wasmUrl telling main.js where the renamed file ended up. In that case, warning when moduleUrl is present without wasmUrl makes sense.

The compile-server path works differently. prepareModel(...) accepts Stan source, asks a compile server to build it, and returns a model whose moduleUrl points to the server-hosted main.js. The documentation says that this main.js loads its matching main.wasm from the same compile server, so there is no local WASM asset for the experiment’s bundler to rename or emit.

The documented use looks like:

const model = await jsPsychADO.prepareModel(sourceSpec, { compileServer });
const ado = jsPsychADO.createController(jsPsych, { model, design_grid });

Now the possible gap...

prepareModel(...) removes the Stan source fields and returns the server-hosted moduleUrl without a wasmUrl. When that result is passed to createController, validateModel can only see a moduleUrl without a wasmUrl, so it emits the warning that a bundled build “would 404 its wasm.” I reproduced that warning with the sequence above.

That warning appears to be applying the local-artifact rule to the server-hosted artifact pair. An author following the documented prepareModel(...) workflow is therefore warned that their configuration is unsafe for bundlers even though this path is documented as loading both artifacts from the compile server.

Could validation retain or recover enough information to distinguish a compile-server result from a locally bundled model, so this workflow receives no warning or a warning appropriate to its actual deployment requirements? A regression test covering prepareModel(...) followed by createController(...) would make the intended behavior explicit.

I realize this behavior appears to predate the PR; I only noticed it while reviewing these changes. Since this PR changes prepareModel validation and develops the compile-server workflow, it seems worth addressing here.

[githubpsyche]

I did a fresh pass over the current head. I think there are three additional issues worth checking. These are separate from the prepareModel(...) warning I raised earlier; the third is a documentation follow-up to the existing readiness discussion.

1. The packed npm package does not include the usage guide that the README now relies on

Reviewing my understanding...

This PR intentionally shortens the root README.md and moves the full API, bundler setup, adaptive-stopping documentation, and related material into docs/usage.md. The README now directs users there for the complete instructions.

The files included in the npm package are controlled by the files array in package.json:

"files": [
  "src",
  "core/tinystan",
  "README.md",
  "CHANGELOG.md",
  "LICENSE"
]

Now the possible gap...

docs/usage.md is not included in that list. I confirmed with npm pack --dry-run that the packed artifact contains README.md but not docs/usage.md.

That means someone reading the README from an installed copy of the package is directed to a local file that does not exist. This matters more after this PR because the full API and bundler guidance previously lived inside the shipped README itself.

Could the package either ship the intended documentation files or use links that remain valid from the packed artifact? Because the usage guide itself links into omitted demo files, it may be worth deciding explicitly whether package documentation should be self-contained or should link to canonical GitHub pages.

2. Source compilation begins before design-grid validation finishes

Reviewing my understanding...

createController(...) validates the supplied model and design grid before returning a usable controller. For a stanCode model, this PR also starts the compile-server request eagerly so compilation can overlap the introductory screens.

The current order in src/index.js is:

if (handle_controller !== "mock") {
  ensureModuleReady();
}

const candidate_designs = enumerateDesigns(config.design_grid);
validateDesignGridForModel(candidate_designs, adapter, adapter.id);

ensureModuleReady() reaches the compile-server fetch(...) immediately. The design-grid validation happens afterward.

Now the possible gap...

A source model with an invalid design grid still sends its Stan source to the compile server before createController(...) throws. I reproduced this with a grid missing a required design key: controller construction failed as expected, but one compile request had already been made.

Could the eager-compilation trigger move below the synchronous model/grid validation? That would preserve essentially all of the intended overlap while ensuring a rejected controller configuration does not perform network work.

3. The public readiness descriptions still reflect the earlier behavior

Reviewing my understanding...

Following the earlier readiness fix, ado.ready() now waits until the Worker has imported the model module and instantiated its WASM. This applies to committed models as well as source-compiled models. For source models, compilation and artifact retrieval happen before that Worker initialization.

Now the possible gap...

docs/usage.md still says that source models wait only for compilation and download while committed-artifact models “resolve immediately.” The ready() documentation in src/index.d.ts says the same thing. src/ado/model_preload.js also describes the promise and participant-facing wait primarily in terms of compilation, although it can now be waiting on a committed model’s Worker/WASM initialization.

Could these descriptions be updated to match the implemented meaning of readiness: Worker import and WASM initialization for both artifact paths, preceded by compile-server work for source models?

I did not find a new failure in the core source-compilation or inference path. The unit, type, browser, and packed-bundler checks all passed; these three points concern package documentation, validation ordering, and keeping the public readiness description aligned with the implementation.