feat(webui): trace SSR runtime dependencies

.gitignore

@@ -187,6 +187,11 @@ webui/frontend/.react-router
 # file it asks for so a local 404-silencer never gets committed again.
 webui/frontend/public/.well-known/appspecific/com.chrome.devtools.json
 
+# The traced runtime tree `pnpm build:image` assembles; `build/` above does not
+# match it. The image builds its own copy inside the WebUI cache rather than here,
+# but a local `pnpm build:image` lands one in the checkout.
+webui/frontend/build-runtime/
+
 # Frontend CSS and dependency caches are generated during build and preparation.
 webui/frontend/public/antd/
 webui/frontend/public/assets/

docker/webui.Dockerfile

@@ -57,8 +57,17 @@ RUN mkdir -p /data /opt/ms-agent-webui-cache
 WORKDIR /app
 # This uses the installed wheel and the final runtime's Node version. No source
 # checkout or frontend compilation is performed inside the image.
+#
+# MS_AGENT_WEBUI_TRACE_RUNTIME keeps only the dependency closure the SSR entries
+# can actually reach: `pnpm install --prod` lands ~430 MB in the cache against a
+# real closure of ~51 MB, the waste sitting inside the packages rather than in a
+# list of unneeded ones. Tracing imports `tsx` and `@vercel/nft`, both
+# devDependencies, so the install it runs on is a full one -- transient, and
+# discarded within this single RUN, so no layer retains it. The tracer boots the
+# closure and renders a page before the full tree goes away, which is why a
+# dependency it could not see fails HERE instead of a request in production.
 RUN --mount=type=cache,target=/root/.local/share/pnpm/store \
-    ms-agent ui --prepare-only --no-browser
+    MS_AGENT_WEBUI_TRACE_RUNTIME=1 ms-agent ui --prepare-only --no-browser
 # Runtime defaults also apply to agent subprocesses with filtered environments.
 # Mount replacement files at these paths to use another package index.
 RUN <<'EOF'

ms_agent/cli/ui.py

@@ -21,6 +21,13 @@
 MIN_UV_VERSION = (0, 5, 0)
 IS_WINDOWS = os.name == 'nt'
 
+#: Opt in to keeping only the traced SSR dependency closure instead of everything
+#: `pnpm install` resolved. Deliberately not a CLI flag: the image sets it while
+#: building, and the first preparation it triggers costs a full devDependency
+#: install plus a tracing pass. Advertising it in `--help` would invite that onto
+#: laptops, where the 430 MB it saves is a cache directory nobody ships.
+TRACE_RUNTIME_ENV = 'MS_AGENT_WEBUI_TRACE_RUNTIME'
+
 
 def _port_number(value):
     try:
@@ -121,20 +128,38 @@ def install_node(frontend, production=False):
 
                 def build_frontend(frontend):
                     pnpm = _require_executable('pnpm')
-                    _check_tool_versions({'node': node, 'pnpm': pnpm}, frontend)
+                    _check_tool_versions({
+                        'node': node,
+                        'pnpm': pnpm
+                    }, frontend)
                     print(
                         '[setup] Building WebUI (including generated CSS)...',
                         flush=True)
                     _run_setup([pnpm, 'build'], frontend, 'frontend build')
 
+                def trace_runtime(frontend):
+                    pnpm = _require_executable('pnpm')
+                    _check_tool_versions({
+                        'node': node,
+                        'pnpm': pnpm
+                    }, frontend)
+                    print(
+                        '[setup] Tracing WebUI runtime dependencies...',
+                        flush=True)
+                    _run_setup(
+                        [pnpm, 'exec', 'tsx', 'scripts/traceRuntime.ts'],
+                        frontend, 'runtime dependency tracing')
+
                 if installed:
+                    tracing = os.environ.get(TRACE_RUNTIME_ENV) == '1'
                     webui = prepare_installed(
                         webui,
                         common,
                         node_version,
                         skip_install=self.args.skip_install,
                         install_node=install_node,
-                        build_frontend=build_frontend)
+                        build_frontend=build_frontend,
+                        trace_runtime=trace_runtime if tracing else None)
                     python = Path(sys.executable)
                 else:
                     backend = webui / 'backend'

ms_agent/cli/ui_resources.py

@@ -128,7 +128,14 @@ def materialize(bundled, manifest, destination):
             shutil.rmtree(temporary)
 
 
-def node_stamp(frontend, node_version, *, production=True):
+def node_stamp(frontend, node_version, *, production=True, traced=False):
+    """Describe the prepared dependency tree, so a stale one is replaced.
+
+    ``production`` describes the wheel's own mode, not the install that produced
+    the tree: tracing needs devDependencies, but that full install is scaffolding
+    for the tracer and is gone before this is written. ``traced`` is what tells a
+    closure apart from everything ``pnpm install`` resolved.
+    """
     return {
         'node':
         list(node_version),
@@ -142,9 +149,51 @@ def node_stamp(frontend, node_version, *, production=True):
         json.loads((frontend / 'package.json').read_text())['packageManager'],
         'production':
         production,
+        'traced':
+        traced,
     }
 
 
+def _stamp_matches(stamp, expected):
+    """Decide whether an already prepared tree satisfies this start.
+
+    Tracing is a property of how the cache was built, not of how it is launched:
+    the image traces once while building and the container then starts from a
+    plain ``CMD`` that carries no such request. A traced tree therefore answers a
+    start that did not ask for one -- otherwise every container start would
+    reject the cache it was shipped with. The reverse must not hold: asking for
+    tracing when the cache holds the full tree has to re-prepare it, or a
+    regression in that step would quietly ship the whole 430 MB again.
+    """
+    if stamp is None:
+        return False
+    if stamp.get('traced') and not expected['traced']:
+        stamp = {**stamp, 'traced': False}
+    return stamp == expected
+
+
+def _prune_to_traced_closure(frontend):
+    """Replace the installed tree with the closure the tracer just verified.
+
+    ``scripts/traceRuntime.ts`` assembles ``build-runtime/`` -- the build output,
+    the entries and only the ``node_modules`` files those entries can reach -- and
+    boots it to render a page before it returns. That order is the point: the full
+    tree is still intact while the closure proves itself, so a dependency the
+    tracer could not see fails before anything has been destroyed.
+    """
+    traced = frontend / 'build-runtime'
+    closure = traced / 'node_modules'
+    if not closure.is_dir():
+        raise UIError('Runtime tracing produced no dependency closure at '
+                      + str(closure))
+    shutil.rmtree(frontend / 'node_modules')
+    # A rename, not a copy: pnpm's layout is relative symlinks into `.pnpm/` and
+    # the tracer recreated them as such, so moving the directory keeps every one
+    # of them resolving inside it.
+    closure.rename(frontend / 'node_modules')
+    shutil.rmtree(traced)
+
+
 def check_backend_dependencies():
     modules = [
         'anthropic', 'exa_py', 'fastapi', 'httpx',
@@ -164,8 +213,14 @@ def check_backend_dependencies():
         )
 
 
-def prepare_installed(bundled, common, node_version, *, skip_install,
-                      install_node, build_frontend=None):
+def prepare_installed(bundled,
+                      common,
+                      node_version,
+                      *,
+                      skip_install,
+                      install_node,
+                      build_frontend=None,
+                      trace_runtime=None):
     from ms_agent.version import __version__
 
     manifest_file = bundled / 'RESOURCE-MANIFEST.json'
@@ -186,20 +241,30 @@ def prepare_installed(bundled, common, node_version, *, skip_install,
         if prebuilt:
             common.validate_build(frontend)
         marker = frontend / '.node-dependencies.json'
-        expected = node_stamp(frontend, node_version, production=prebuilt)
+        expected = node_stamp(
+            frontend,
+            node_version,
+            production=prebuilt,
+            traced=bool(trace_runtime))
         try:
-            ready = json.loads(marker.read_text()) == expected and (
-                frontend / 'node_modules').is_dir()
+            stamp = json.loads(marker.read_text())
         except (OSError, ValueError):
-            ready = False
+            stamp = None
+        ready = _stamp_matches(
+            stamp, expected) and (frontend / 'node_modules').is_dir()
         if not ready:
             if skip_install:
                 raise UIError(
                     'WebUI Node dependencies need preparation. Run once without --skip-install.'
                 )
             marker.unlink(missing_ok=True)
-            install_node(frontend, production=prebuilt)
-            marker.write_text(json.dumps(expected, sort_keys=True) + '\n')
+            # The tracer runs out of this tree and imports `tsx` and
+            # `@vercel/nft` from it, so tracing cannot start from a production
+            # install. What it leaves behind is narrower than either.
+            install_node(
+                frontend, production=prebuilt and trace_runtime is None)
+            if trace_runtime is None:
+                marker.write_text(json.dumps(expected, sort_keys=True) + '\n')
         if not prebuilt:
             try:
                 common.validate_build(frontend)
@@ -212,4 +277,11 @@ def prepare_installed(bundled, common, node_version, *, skip_install,
                     raise UIError('No WebUI frontend builder is available')
                 build_frontend(frontend)
                 common.validate_build(frontend)
+        if not ready and trace_runtime is not None:
+            # After any build, because the tracer follows what the SSR bundle
+            # imports. The stamp lands only once the closure is in place: an
+            # interrupted trace must not leave one claiming this tree is traced.
+            trace_runtime(frontend)
+            _prune_to_traced_closure(frontend)
+            marker.write_text(json.dumps(expected, sort_keys=True) + '\n')
     return destination

tests/cli/test_ui.py

@@ -261,3 +261,68 @@ def test_reload_fails_with_development_instructions(capsys):
         ui.UICMD(_parse_ui_args('--reload')).execute()
     assert error.value.code == 1
     assert 'pnpm dev' in capsys.readouterr().err
+
+
+def _prepare_installed_kwargs(monkeypatch, trace_env=None):
+    """Drive execute() far enough to capture what it hands prepare_installed.
+
+    `--prepare-only` returns before ports and the launcher, so the WebUI backend
+    never has to be importable here: a stub `app.processes` satisfies the single
+    import execute() makes for its own signal handling.
+    """
+    import contextlib
+    import sys
+    import types
+    from pathlib import Path
+
+    package = types.ModuleType('app')
+    package.__path__ = []
+
+    @contextlib.contextmanager
+    def interruptible():
+        yield
+
+    processes = types.ModuleType('app.processes')
+    processes.interruptible = interruptible
+    monkeypatch.setitem(sys.modules, 'app', package)
+    monkeypatch.setitem(sys.modules, 'app.processes', processes)
+
+    if trace_env is None:
+        monkeypatch.delenv(ui.TRACE_RUNTIME_ENV, raising=False)
+    else:
+        monkeypatch.setenv(ui.TRACE_RUNTIME_ENV, trace_env)
+
+    captured = {}
+    monkeypatch.setattr(ui, 'find_webui', lambda: (Path('/webui'), True))
+    monkeypatch.setattr(ui, 'load_common', lambda webui: SimpleNamespace())
+    monkeypatch.setattr(ui, '_require_executable', lambda name: '/tools/' + name)
+    monkeypatch.setattr(ui, '_check_tool_versions',
+                        lambda tools, frontend=None: None)
+    monkeypatch.setattr(ui, '_read_semantic_version',
+                        lambda *args, **kwargs: (22, 22, 0))
+
+    def fake_prepare(webui, common, node_version, **kwargs):
+        captured.update(kwargs)
+        return webui
+
+    monkeypatch.setattr(ui, 'prepare_installed', fake_prepare)
+    ui.UICMD(_parse_ui_args('--prepare-only', '--no-browser')).execute()
+    return captured
+
+
+def test_runtime_tracing_is_off_unless_the_environment_asks(monkeypatch):
+    assert _prepare_installed_kwargs(monkeypatch)['trace_runtime'] is None
+
+
+def test_runtime_tracing_env_reaches_preparation(monkeypatch):
+    kwargs = _prepare_installed_kwargs(monkeypatch, trace_env='1')
+    assert callable(kwargs['trace_runtime'])
+
+
+@pytest.mark.parametrize('value', ['0', '', 'true', 'yes', 'false'])
+def test_only_an_exact_1_enables_runtime_tracing(monkeypatch, value):
+    """A truthiness check here would trace on `=0`, and the cost of that (a full
+    devDependency install plus a tracing pass) lands on whoever set the variable
+    specifically to turn the feature off."""
+    kwargs = _prepare_installed_kwargs(monkeypatch, trace_env=value)
+    assert kwargs['trace_runtime'] is None

tests/cli/test_ui_resources.py

@@ -169,3 +169,94 @@ def fail(frontend):
         resources.prepare_installed(root, common, (22, 23, 2), skip_install=False,
                                     install_node=lambda *a, **kw: pytest.fail('must not install'),
                                     build_frontend=lambda *a: pytest.fail('must not rebuild'))
+
+
+@pytest.fixture
+def traced(bundle, tmp_path, monkeypatch):
+    """A prebuilt cache whose Node preparation can be traced down to its closure."""
+    root, _ = bundle
+    monkeypatch.setenv('MS_AGENT_WEBUI_CACHE', str(tmp_path / 'runtime'))
+    monkeypatch.setattr(resources, 'check_backend_dependencies', lambda: None)
+    return SimpleNamespace(
+        bundled=root,
+        common=SimpleNamespace(validate_build=lambda directory: None),
+        cache=tmp_path / 'runtime')
+
+
+def _install_everything(frontend, production):
+    """Stand in for `pnpm install`: a tree carrying a file no closure can reach."""
+    modules = frontend / 'node_modules'
+    modules.mkdir()
+    (modules / 'unreachable.js').write_text('// the 430 MB nobody loads')
+
+
+def _fake_trace(frontend):
+    """Stand in for traceRuntime.ts: assemble build-runtime/ the way it does,
+    including pnpm's relative symlink layout, which is what the move must keep."""
+    closure = frontend / 'build-runtime/node_modules'
+    (closure / '.pnpm/antd@6/node_modules/antd').mkdir(parents=True)
+    (closure / '.pnpm/antd@6/node_modules/antd/index.js').write_text('// SSR')
+    (closure / 'antd').symlink_to('.pnpm/antd@6/node_modules/antd')
+    (frontend / 'build-runtime/server.js').write_text('// entry')
+
+
+def test_tracing_replaces_the_installed_tree_with_its_closure(traced):
+    installs = []
+
+    def install(frontend, production):
+        # Tracing imports tsx and @vercel/nft out of this very tree, so it can
+        # never be a production install.
+        assert production is False
+        installs.append(frontend)
+        _install_everything(frontend, production)
+
+    prepared = resources.prepare_installed(traced.bundled, traced.common, (22, 23, 2),
+                                           skip_install=False, install_node=install,
+                                           trace_runtime=_fake_trace)
+    modules = prepared / 'frontend/node_modules'
+    assert len(installs) == 1
+    assert not (modules / 'unreachable.js').exists()
+    # Moved, not copied: the relative symlink still resolves inside the tree.
+    assert (modules / 'antd').is_symlink()
+    assert (modules / 'antd/index.js').read_text() == '// SSR'
+    assert not (prepared / 'frontend/build-runtime').exists()
+
+
+def test_traced_cache_is_reused_by_a_plain_start(traced):
+    prepared = resources.prepare_installed(traced.bundled, traced.common, (22, 23, 2),
+                                           skip_install=False, install_node=_install_everything,
+                                           trace_runtime=_fake_trace)
+    # Exactly what the container does: CMD carries --skip-install and no tracing
+    # request, because tracing already happened while the image was built. This
+    # start must accept the tree it was shipped, not reject it as unprepared.
+    reused = resources.prepare_installed(traced.bundled, traced.common, (22, 23, 2),
+                                         skip_install=True,
+                                         install_node=lambda *a, **kw: pytest.fail('must not install'))
+    assert reused == prepared
+
+
+def test_untraced_cache_is_re_prepared_when_tracing_is_requested(traced):
+    resources.prepare_installed(traced.bundled, traced.common, (22, 23, 2),
+                                skip_install=False, install_node=_install_everything)
+    # The reverse of the case above must NOT hold: accepting the full tree here
+    # is how the image would quietly go back to shipping 430 MB.
+    with pytest.raises(resources.UIError, match='need preparation'):
+        resources.prepare_installed(traced.bundled, traced.common, (22, 23, 2),
+                                    skip_install=True,
+                                    install_node=lambda *a, **kw: pytest.fail('must not install'),
+                                    trace_runtime=lambda frontend: pytest.fail('must not trace'))
+
+
+def test_tracing_that_produced_no_closure_keeps_the_installed_tree(traced):
+    def trace_without_closure(frontend):
+        (frontend / 'build-runtime').mkdir()
+
+    with pytest.raises(resources.UIError, match='no dependency closure'):
+        resources.prepare_installed(traced.bundled, traced.common, (22, 23, 2),
+                                    skip_install=False, install_node=_install_everything,
+                                    trace_runtime=trace_without_closure)
+    frontend = next(traced.cache.glob('*/frontend'))
+    # Nothing half-pruned, and no stamp claiming this tree is a closure: the next
+    # attempt has to find it unprepared rather than serve an incomplete tree.
+    assert (frontend / 'node_modules/unreachable.js').exists()
+    assert not (frontend / '.node-dependencies.json').exists()