mlx-bun 0.2.0 → 0.3.0

package/src/model/trellis-wide-prefill.ts

@@ -0,0 +1,94 @@
+// Packed Trellis prefill with MLX 0.32.2 GemvWide arithmetic.
+// Ported from ml-explore/mlx v0.32.2, mlx/backend/metal/kernels/gemv.h
+// (MIT). Decode weights to bf16 inside the native dot/unroll/shuffle order;
+// share each decoded row across every input vector in this narrow regime.
+import { MetalKernel } from "../mlx/metal-kernel";
+import { type MlxArray } from "../mlx/array";
+import { Dtype, deviceArchitecture } from "../mlx/ffi";
+import { type TrellisGeometry } from "./trellis-linear";
+
+/** Matches the bundled runtime's native reduction on M3 and newer GPUs. */
+export function nativeTrellisWidePrefill(m: number): boolean {
+  return m >= 5 && m <= 15 &&
+    Number(/^applegpu_[a-z](\d{2})/.exec(deviceArchitecture())?.[1] ?? 0) >= 15;
+}
+
+export function wideTrellisPrefillEligible(g: TrellisGeometry, m: number, dtype: Dtype): boolean {
+  return dtype === Dtype.bfloat16 && nativeTrellisWidePrefill(m) &&
+    g.axis === 1 && g.T === 256 && g.L === 12 && [2, 3, 4].includes(g.k) &&
+    g.inFeatures === 5120 && g.outFeatures === 17408 && !g.blockInterleave;
+}
+
+let kernel: MetalKernel | undefined;
+function wideKernel() {
+  return kernel ??= new MetalKernel({ name: 'mlx_bun_trellis_wide_prefill',
+    inputNames: ['x', 'codes', 'scales'], outputNames: ['out'], source: String.raw`
+    constexpr uint unroll=8, rows_per_simd=32/KL, simdgroups=KL/8;
+    const uint lane=thread_index_in_simdgroup, sg=simdgroup_index_in_threadgroup;
+    const uint kl=lane%KL, row=threadgroup_position_in_grid.y*4+sg*rows_per_simd+lane/KL;
+    const uint vec0=threadgroup_position_in_grid.x*SM;
+    constexpr uint wpb=BT*BITS/32, row_words=KC*BITS/32;
+    const float scale=float(scales[row]);
+    constexpr uint n_v4=KC/4, n_main=n_v4-n_v4%(KL*unroll);
+    float result[SM]={0};
+    for(uint base=0;base<n_main;base+=KL*unroll) {
+      float4 wf[unroll];
+      #pragma clang loop unroll(full)
+      for(uint i=0;i<unroll;i++) {
+        #pragma clang loop unroll(full)
+        for(uint c=0;c<4;c++) {
+          const uint k=(base+i*KL+kl)*4+c, block=k/BT, t=k%BT;
+          const device uint* bw=codes+(ulong)row*row_words+block*wpb;
+          const uint pos=(BT-1-t)*BITS, wi=pos>>5, off=pos&31u;
+          uint win=bw[wi]>>off;
+          if(off+L>32)win|=bw[wi+1==wpb?0:wi+1]<<(32-off);
+          win&=(1u<<L)-1;
+          const uint z=win*34038481u+76625530u;
+          const uint p=(z&0x00FF00FFu)+((z>>8)&0x00FF00FFu);
+          const int code=int((p&0xFFFFu)+(p>>16))-510;
+          const float reciprocal=1.0f/147.800537109375f;
+          const float approx=float(code)*reciprocal;
+          const float residual=metal::fma(-approx,147.800537109375f,float(code));
+          const float value=metal::fma(residual,reciprocal,approx);
+          wf[i][c]=float(T(value*scale));
+        }
+      }
+      #pragma clang loop unroll(full)
+      for(uint v=0;v<SM;v++) {
+        const uint m=min(vec0+v,uint(M-1));
+        float acc=0;
+        #pragma clang loop unroll(full)
+        for(uint i=0;i<unroll;i++) {
+          const ulong offset=(ulong)m*KC+(base+i*KL+kl)*4;
+          const float4 xv=float4(float(x[offset]),float(x[offset+1]),float(x[offset+2]),float(x[offset+3]));
+          acc+=metal::dot(wf[i],xv);
+        }
+        result[v]+=acc;
+      }
+    }
+    #pragma clang loop unroll(full)
+    for(uint v=0;v<SM;v++) {
+      #pragma clang loop unroll(full)
+      for(ushort off=KL/2;off>=1;off>>=1) result[v]+=simd_shuffle_down(result[v],off);
+      if(kl==0 && vec0+v<M) out[(ulong)(vec0+v)*N+row]=T(result[v]);
+    }
+  ` });
+}
+/** Requires the caller to prove aligned, row-contiguous x, as Qwen does
+ * with its RMSNorm output. Preserve native fallback when that proof is absent. */
+export function wideTrellisPrefill(x: MlxArray, codes: MlxArray, scales: MlxArray, g: TrellisGeometry): MlxArray {
+  const [m, k] = x.shape;
+  if (x.ndim !== 2 || x.dtype !== Dtype.bfloat16 || !m || m < 5 || m > 15 ||
+      k !== 5120 || g.inFeatures !== k || g.axis !== 1 || g.T !== 256 || g.L !== 12 ||
+      ![2, 3, 4].includes(g.k) || g.blockInterleave || g.rows < 4 || g.rows % 4 ||
+      g.outFeatures !== g.rows || g.cols !== k || codes.dtype !== Dtype.uint32 ||
+      codes.ndim !== 2 || codes.shape[0] !== g.rows || codes.shape[1] !== k * g.k / 32 || scales.size !== g.rows)
+    throw new Error("unsupported wide trellis prefill geometry, dtype or layout");
+  const kl = m <= 5 ? 32 : 16;
+  return wideKernel().apply([x, codes, scales], {
+    outputs: [{ shape: [m, g.outFeatures], dtype: x.dtype }],
+    grid: [32, g.outFeatures / 4 * (kl / 8), 1], threadGroup: [32, kl / 8, 1],
+    templateDtypes: { T: x.dtype },
+    templateInts: { M: m, SM: m, KL: kl, KC: k, N: g.outFeatures, BT: g.T, L: g.L, BITS: g.k },
+  })[0]!;
+}

package/src/model/universal/dense.ts

@@ -3,7 +3,7 @@
 // like qwen3.ts, with every arch delta selected by the UniversalArgs
 // descriptor (archs.ts). Monolith path only: no compiled
 // decode, no kv-quant — slow, never broken. The bar is L1 bit-exactness
-// vs mlx-lm on this machine's GPU (tests/universal-parity.test.ts).
+// vs mlx-lm on this machine's GPU (tests/parity/universal-parity.test.ts).
 //
 // Porting discipline: each branch transcribes its mlx-lm source op-for-op
 // (q/k-norm before vs after reshape, rope order, scale points, softcap,

package/src/model/universal/modules.ts

@@ -11,6 +11,7 @@ import type { Weights } from "../../weights";
 import { MlxArray } from "../../mlx/array";
 import * as ops from "../../mlx/ops";
 import { QuantizedEmbedding, QuantizedLinear, RMSNorm } from "../gemma4-base";
+import { TrellisLinear } from "../trellis-linear";
 
 /** Records every tensor a universal load consumes; `finish` diffs against
  *  the shard index so a descriptor mistake is a LOAD error (unconsumed /
@@ -66,7 +67,7 @@ export class DenseLinear {
 }
 
 /** Either linear flavor behind one call surface. */
-export type AnyLinear = QuantizedLinear | DenseLinear;
+export type AnyLinear = QuantizedLinear | DenseLinear | TrellisLinear;
 
 /** Quantized when `.scales` exists (MLX-quantized checkpoints), else dense
  *  bf16/f16/f32 (Phase 1.5). Both flavors carry the optional ADDITIVE
@@ -74,6 +75,12 @@ export type AnyLinear = QuantizedLinear | DenseLinear;
 export function loadLinear(
   weights: Weights, path: string, config: ModelConfig, audit: WeightAudit,
 ): AnyLinear {
+  if (TrellisLinear.isTrellis(config, path)) {
+    const lin = TrellisLinear.load(weights, path, config);
+    audit.use(`${path}.weight`);
+    audit.use(`${path}.scales`);
+    return lin;
+  }
   if (weights.has(`${path}.scales`)) {
     const lin = QuantizedLinear.load(weights, path, config);
     audit.use(`${path}.weight`);

package/src/native-pack.ts

@@ -8,7 +8,7 @@
 //
 // Resolution order for the libs (mirrored in src/mlx/ffi.ts):
 //   MLX_BUN_LIBMLXC env > beside the executable (sidecar dist/ layout,
-//   docs/reference/embedding.md) > the native-pack cache dir > homebrew.
+//   docs/reference/distribution.md) > the native-pack cache dir > homebrew.
 // Dev trees and embedders therefore never download anything.
 
 import { existsSync, mkdirSync, renameSync, rmSync } from "node:fs";
@@ -17,22 +17,23 @@ import { dirname, join } from "node:path";
 import { downloadOne } from "./download";
 import { runtimeValue } from "./runtime-config";
 
-export const NATIVE_PACK_VERSION = "0.3.0";
+export const NATIVE_PACK_VERSION = "0.4.0";
+export const MLX_CORE_VERSION = "0.32.2";
 export const NATIVE_PACK_FILES = [
   "libmlxc.dylib", "libmlx.dylib", "libjaccl.dylib", "mlx.metallib",
   "libmlx_bun_expert_io.dylib", "mlx-bun-frame-extract",
 ] as const;
 
 // RELEASE-TIME BAKE: scripts/build-native-pack.sh prints the sha256/size for
-// the freshly built tarball — publish `native-v0.3.0` with both assets and
+// the freshly built tarball — publish `native-v0.4.0` with both assets and
 // paste the values here BEFORE tagging the package release (the download
 // 404s until the GitHub release exists; dev trees resolve via homebrew/env
 // and never hit this path).
 const SHA256: Record<string, string> = {
-  arm64: "b485fed858971ea6f072f349a0234c0e2d5030d3ef2ddfa942c8975eca2d3b2e",
+  arm64: "f2b0857658aa0826aee14b94a6b83d4c86c08f2c07d716991de6afd038f5ca5d",
 };
 const SIZE: Record<string, number> = {
-  arm64: 52299265,
+  arm64: 41931583,
 };
 
 export function nativePackName(arch = process.arch): string {

package/src/pi-provider.ts

@@ -17,37 +17,11 @@ import { AuthStorage, ModelRegistry } from "@earendil-works/pi-coding-agent";
  *  derive it from the registry's `find()` return type. */
 export type PiModel = NonNullable<ReturnType<ModelRegistry["find"]>>;
 
-// Provider/model constants — kept aligned with src/harness-pi.ts so every
-// path (web embed, terminal embed, external pi) registers the same provider.
-export const PI_PROVIDER_ID = "mlx-bun";
-export const PI_LOCAL_MODEL_ID = "local";
-export const PI_API_KEY = "sk-mlx-bun-local";
-export const PI_API = "openai-completions" as const;
-
-export const DEFAULT_CONTEXT_WINDOW = 32_768;
-export const DEFAULT_MAX_TOKENS = 8_192;
-
-export interface PiProviderOptions {
-  /** Advertised context window. Default: DEFAULT_CONTEXT_WINDOW. */
-  contextWindow?: number;
-  /** Advertised max output tokens. Default: DEFAULT_MAX_TOKENS. */
-  maxTokens?: number;
-  /** Human-readable model label (e.g. the real repo id). Default: "mlx-bun (local)". */
-  name?: string;
-  /** The loaded model has a switchable reasoning channel (server's
-   *  `ctx.template.supportsThinking` — true for any model whose chat template
-   *  gates on enable_thinking: Qwen3.5, MiniCPM5, …). When true Pi engages
-   *  reasoning and sends the on/off switch the server maps to enable_thinking.
-   *  Default false. */
-  reasoning?: boolean;
-  /** The loaded model can accept images (server's `ctx.vision != null`).
-   *  This MUST match the real model: pi-ai serializes user content as OpenAI
-   *  multimodal content-parts (`[{type:"text",…}]`) for image-capable models
-   *  and as a plain string otherwise. A text-only chat template renders nothing
-   *  for a parts array, so declaring a non-vision model as image-capable makes
-   *  the user's turn vanish ("I don't see any message"). Default false. */
-  vision?: boolean;
-}
+import { PI_PROVIDER_ID, PI_API_KEY, PI_API, PI_LOCAL_MODEL_ID, piModelDefinition,
+  type PiModelOptions } from "./contracts/pi-provider";
+export { PI_PROVIDER_ID, PI_API_KEY, PI_API, PI_LOCAL_MODEL_ID,
+  DEFAULT_CONTEXT_WINDOW, DEFAULT_MAX_TOKENS } from "./contracts/pi-provider";
+export interface PiProviderOptions extends PiModelOptions {}
 
 export interface PiProviderWiring {
   authStorage: AuthStorage;
@@ -70,39 +44,7 @@ export function buildPiProvider(baseUrl: string, opts: PiProviderOptions = {}):
     baseUrl,
     apiKey: PI_API_KEY,
     api: PI_API,
-    models: [
-      {
-        id: PI_LOCAL_MODEL_ID,
-        name: opts.name ?? "mlx-bun (local)",
-        api: PI_API,
-        reasoning: opts.reasoning ?? false,
-        // Compat flags that make pi emit the wire format our local server's
-        // chat templates expect (docs/custom-provider.md "Model Definition
-        // Reference"):
-        //  - supportsDeveloperRole:false → pi sends the system prompt as the
-        //    `system` role, not `developer` (our templates only know
-        //    system/user/assistant/tool).
-        //  - thinkingFormat:"qwen-chat-template" (reasoning models only) → pi
-        //    sends chat_template_kwargs.enable_thinking, the boolean our
-        //    Qwen3.5/MiniCPM5 templates gate on. pi computes the boolean from
-        //    the selected thinking level (off → false), so no thinkingLevelMap.
-        compat: {
-          supportsDeveloperRole: false,
-          ...(opts.reasoning ? { thinkingFormat: "qwen-chat-template" as const } : {}),
-        },
-        // input MUST match the loaded model. pi-ai serializes user content as a
-        // plain string for text-only models and as OpenAI content-parts
-        // (`[{type:"text",…}]`) once a model declares image input — and a
-        // text-only chat template renders NOTHING for a parts array, dropping
-        // the user's turn ("I don't see any message"). So declare image input
-        // only for an actual vision model; the web `ready` frame's `vision`
-        // flag is the same signal that drives the UI's image affordance.
-        input: opts.vision ? ["text", "image"] : ["text"],
-        cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
-        contextWindow: opts.contextWindow ?? DEFAULT_CONTEXT_WINDOW,
-        maxTokens: opts.maxTokens ?? DEFAULT_MAX_TOKENS,
-      },
-    ],
+    models: [piModelDefinition(opts)],
   });
 
   const model = modelRegistry.find(PI_PROVIDER_ID, PI_LOCAL_MODEL_ID);

package/src/pi-web.ts

@@ -1,3 +1,7 @@
+import { APP_ROUTE_IDS } from "./contracts/pi";
+export { APP_ROUTE_IDS } from "./contracts/pi";
+import type { ApprovalDecision, ImageAttachment, ClientMessage, UiSnapshotElement, UiSnapshot, WizardStep, AppUiContext, AppRouteId, ReadyGenDefaults, ServerMessage, HistoryToolItem, HistoryItem, SessionListItem, SamplingOverrides, Lane } from "./contracts/pi";
+export type { ApprovalDecision, ImageAttachment, ClientMessage, UiSnapshotElement, UiSnapshot, WizardStep, AppUiContext, AppRouteId, ReadyGenDefaults, ServerMessage, HistoryToolItem, HistoryItem, SessionListItem, SamplingOverrides } from "./contracts/pi";
 // mlx-bun pi-web — server-side embed of the pi agent SDK that bridges a
 // browser WebSocket to an in-process AgentSession driving mlx-bun's own
 // local model (Phase 16, web half).
@@ -52,7 +56,7 @@ import { MEMORY_TOOL_NAMES, REFERENCE_TOOL_NAMES } from "./memory/tools";
 import { buildPiAgentSurface } from "./pi-session";
 import { buildPiProvider, DEFAULT_CONTEXT_WINDOW, PI_LOCAL_MODEL_ID } from "./pi-provider";
 import { downloadsSnapshot } from "./download";
-import { getLane, type Lane } from "./serve/lane-registry";
+import { getLane } from "./serve/lane-registry";
 import { isToolAlwaysAllowed, setToolAlwaysAllowed, listAlwaysAllowedTools } from "./tool-approvals";
 
 /**
@@ -356,188 +360,7 @@ export interface PiWsData {
 
 type PiWs = ServerWebSocket<PiWsData>;
 
-// ---- WS protocol message shapes --------------------------------------
-
-/** Decision a browser returns for a gated tool call. `editedArgs`, when
- *  present on an "allow", replaces the tool's proposed input before it runs
- *  (LM Studio's editable-args pattern — plan §5.4/§6.5). `alwaysAllow` marks
- *  the decision as "and remember this tool" — handled by resolveApproval,
- *  which persists it to the durable tool-approvals config (src/tool-
- *  approvals.ts) before settling the pending promise. */
-export type ApprovalDecision = "allow" | "deny";
-
-/** An image attachment from the browser: base64 (no data: prefix) + mime. */
-export interface ImageAttachment {
-  data: string;
-  mimeType: string;
-}
-
-/** Client -> server frames. Exported type-only for the web frontend
- *  (src/web/src/*.ts import it via `import type` — see tsconfig.web.json
- *  and scripts/build-web.ts) so the WS contract is checked at compile time
- *  on both ends instead of drifting between server TS and untyped inline
- *  browser JS (the observed Phase-0 bug class). */
-export type ClientMessage =
-  | { type: "prompt"; text: string; images?: ImageAttachment[] }
-  | { type: "abort" }
-  // `editedArgs`: when the user edited the proposed arguments in the
-  // approval card's textarea before approving, the (re-parsed) JSON object
-  // to substitute for the tool's original input — mutated onto the pi SDK's
-  // ToolCallEvent.input in place (the SDK's documented mutation contract;
-  // see installApprovalGate). Ignored on "deny". `alwaysAllow`: persist this
-  // tool to the durable always-allow config (src/tool-approvals.ts) so
-  // future calls to the SAME TOOL NAME skip the card entirely, on this and
-  // every other session.
-  | { type: "approval"; callId: string; decision: ApprovalDecision; editedArgs?: Record<string, unknown>; alwaysAllow?: boolean }
-  // Toggle the model's reasoning channel on/off (only meaningful when the
-  // model supports thinking — ready.thinking). Maps to Pi's session thinking
-  // level: "medium" (on) ↔ "off". Pi sends it as enable_thinking to the server.
-  | { type: "set_thinking"; enabled: boolean }
-  // Opt-in toggle (plan §5.4/§6.5/§9 Phase 2, default OFF): "let the agent
-  // touch files on this machine". Persisted per-browser by the client
-  // (localStorage `mlxbun.codingTools`); the ALLOWLIST decision is enforced
-  // here, server-side. buildPiAgentSurface's tool list is fixed for the
-  // life of an AgentSession (pi's DefaultResourceLoader config is baked in
-  // at createAgentSessionServices time), so this can't rewire an
-  // in-progress session — it takes effect starting with the next
-  // new_session/open_session/fork_session (createRuntimeFactory reads
-  // this.codingToolsRequested fresh every time the SDK calls it, including
-  // from AgentSessionRuntime.newSession()'s internal re-invocation). The
-  // `coding_tools` ServerMessage below reports the honest current-vs-pending
-  // state so the UI can say so instead of implying it applies immediately.
-  | { type: "set_coding_tools"; enabled: boolean }
-  // Select the active LoRA adapter for subsequent turns (null = none/base).
-  // app.html mounts it (POST /v1/adapters) before sending this; the
-  // before_provider_request hook injects it into the provider payload.
-  | { type: "set_adapter"; id: string | null }
-  // Per-request sampling overrides. Each field is optional; null/undefined
-  // means "leave it to the server default" (the mode-aware recommended
-  // value resolved in toOptions). A present numeric value is injected into
-  // the provider payload and always wins. The before_provider_request hook
-  // injects whatever is set here.
-  //
-  // `scope` (plan §9 Phase 3 / beat matrix Axis 4 "per-message sampling
-  // scope"): "session" (default, omitting the field means this too) persists
-  // the override on the connection until changed — today's behavior,
-  // unchanged. "next_turn" stores a ONE-SHOT override applied to exactly
-  // the next prompt's injection, then cleared automatically — it never
-  // touches the session-level values (a later plain set_sampling with
-  // scope:"session" or omitted still sees whatever was there before this
-  // message). Composition when both are present: next_turn's SET fields
-  // win over session's for that one request; anything next_turn leaves
-  // null/unset falls back to the session-level override for that field
-  // (composeSampling — session-level itself falls back to the server's
-  // mode-aware default per injectSampling's existing contract).
-  | {
-      type: "set_sampling";
-      scope?: "session" | "next_turn";
-      temperature?: number | null;
-      top_p?: number | null;
-      top_k?: number | null;
-      min_p?: number | null;
-      xtc_probability?: number | null;
-      xtc_threshold?: number | null;
-      repetition_penalty?: number | null;
-      repetition_context_size?: number | null;
-      presence_penalty?: number | null;
-      frequency_penalty?: number | null;
-      seed?: number | null;
-    }
-  // Per-session USER system prompt, layered ON TOP OF (not replacing) the
-  // built-in surface prompt (buildWebChatSystemPrompt + memoryHint). null
-  // clears it. Applied by the before_agent_start hook on the NEXT turn (pi
-  // calls it fresh inside every session.prompt(), so this is a true
-  // per-turn injection, not a session-creation-only setting — see
-  // installSystemPromptHook). Presets v1 (composer.ts) send this alongside
-  // set_sampling to apply a saved bundle.
-  | { type: "set_system_prompt"; text: string | null }
-  // Session management (recent-chats sidebar + new chat).
-  | { type: "new_session" }
-  | { type: "list_sessions" }
-  | { type: "open_session"; path: string }
-  | { type: "fork_session"; path: string }
-  | { type: "delete_session"; path: string }
-  // Message actions (plan §5.2). All three navigate the session's leaf
-  // pointer via AgentSession.navigateTree() then re-prompt — a sibling
-  // branch in the SAME session file, never a new one (see the helpers above
-  // findLastUserMessageEntry for the full rationale). Regenerate re-sends the
-  // last user message's original content verbatim; edit_resend re-sends it
-  // with edited text (images preserved); switch_sibling moves the active
-  // leaf to a previously-created sibling (an earlier edit/regeneration)
-  // without sending anything new.
-  | { type: "regenerate" }
-  | { type: "edit_resend"; text: string }
-  | { type: "switch_sibling"; entryId: string }
-  // App-aware assistant (plan §6.6, §9 Phase 3, beat matrix Axis 12): pushed
-  // by chat.ts on every route change AND on wizard-step change (quantize/
-  // finetune/dataset), never on a timer. `context` is opaque here (the
-  // browser's AppContext shape, src/web/src/assistant.ts) — stored verbatim
-  // as the connection's currentAppContext and returned by the
-  // get_current_app_context tool, plus mined for the compact one-line
-  // ambient context auto-prepended to the NEXT prompt (see
-  // installAppContextHook / ambientContextLine). This is a DIFFERENT frame
-  // from the server->client `context` ServerMessage (context-window usage,
-  // above) — same name, opposite direction, disambiguated by which type
-  // union it's a member of; the two are never in scope at the same call
-  // site (ClientMessage here vs. ServerMessage below).
-  | { type: "context"; context: AppUiContext };
-
-// ---- App-aware assistant (plan §6.6, §9 Phase 3, beat matrix Axis 12) --
-//
-// Mirrors src/web/src/assistant.ts's browser-side types (which are the
-// source of truth for shape — this server-side copy is intentionally loose
-// (mostly `unknown`/optional) since the server never constructs a
-// UiSnapshot itself, only stores and forwards what the browser sends. Kept
-// here (not imported) because src/web/src/assistant.ts is bundled into the
-// browser and importing it here would pull browser code into the server
-// build — the inverse of the type-only ClientMessage/ServerMessage import
-// direction the frontend uses for pi-web.ts.
-
-/** One element in a captured DOM snapshot — see assistant.ts's
- *  UiSnapshotElement for the authoritative shape/doc. */
-export interface UiSnapshotElement {
-  ref: string;
-  tag: string;
-  label: string;
-  kind: "interactive" | "region";
-  role?: string;
-  selector: string;
-  spotlightId?: string;
-}
-
-export interface UiSnapshot {
-  route: string;
-  capturedAt: string;
-  elements: UiSnapshotElement[];
-}
-
-/** The wizard-step process-state fields (§6.6 "carries process state, not
- *  just elements") — present only on routes with a step wizard. */
-export interface WizardStep {
-  index: number;
-  count: number;
-  label: string;
-}
-
-/** The full context payload pushed by chat.ts on route/step change — see
- *  assistant.ts's buildAppContext. Stored verbatim per-connection. */
-export interface AppUiContext {
-  route: string;
-  view?: string;
-  step?: WizardStep;
-  snapshot: UiSnapshot;
-}
-
-/** Route catalog mirror (src/web/src/ui-catalog.ts's ROUTE_IDS) for
- *  server-side validation of navigate_app's route/page param BEFORE it
- *  ever reaches the browser — an unknown route is rejected here as a tool
- *  error rather than sent as a `ui_navigate` frame the router can't match.
- *  Keep in sync with ui-catalog.ts's ROUTE_IDS by hand (no shared-JSON
- *  import between the server and browser bundles — see that file's
- *  top-of-file comment for why). */
-export const APP_ROUTE_IDS = ["chat", "quantize", "finetune", "dataset", "status"] as const;
-export type AppRouteId = (typeof APP_ROUTE_IDS)[number];
-
+/** Validate assistant navigation against the shared browser/server route IDs. */
 export function isAppRouteId(v: string): v is AppRouteId {
   return (APP_ROUTE_IDS as readonly string[]).includes(v);
 }
@@ -570,117 +393,6 @@ export function ambientContextLine(ctx: AppUiContext | null | undefined): string
   return `[user is on: ${place}${stepPart}]`;
 }
 
-/** Model-author recommended sampling (generation_config.json, server-CLI
- *  overrides applied) — sent once on `ready` so the sampling popover can show
- *  per-model defaults instead of one hardcoded shape (web-ui-pass-plan.md
- *  #14). A field is null when neither the model nor the server configured one
- *  (the built-in fallback in toOptions still applies server-side). */
-export interface ReadyGenDefaults {
-  temperature: number | null;
-  topP: number | null;
-  topK: number | null;
-}
-
-/** Server -> client frames. Exported type-only — see ClientMessage above. */
-export type ServerMessage =
-  // `vision`: whether the loaded model can accept images (drives the UI's
-  // image-attach affordance — false on e4b until the SigLIP sidecar lands).
-  // `audio`: whether the loaded model can accept `input_audio` parts (audio
-  // tower loaded or loadable). Discovery-only for now: pi's Model.input type
-  // has no "audio" modality, so unlike `vision` this flag never feeds the
-  // provider declaration — clients use it instead of probing for a 400.
-  // `thinking`: whether the model has a switchable reasoning channel (drives
-  // the UI's thinking on/off toggle; false hides it).
-  | { type: "ready"; model: string; vision: boolean; audio: boolean; thinking: boolean; genDefaults: ReadyGenDefaults }
-  | { type: "turn_start" }
-  | { type: "text_delta"; delta: string }
-  | { type: "thinking_delta"; delta: string }
-  | { type: "tool_start"; callId: string; tool: string; args: unknown }
-  | { type: "tool_approval_request"; callId: string; tool: string; args: unknown }
-  | { type: "tool_update"; callId: string; chunk: unknown }
-  | { type: "tool_end"; callId: string; ok: boolean; result: unknown }
-  // `lane`: the serving lane the just-finished turn ran on (serial /
-  // serial+spec / batched), correlated via the lane registry keyed by the
-  // AssistantMessage's responseId (the pi SDK's own usage parsing drops
-  // custom fields — see src/serve/lane-registry.ts). Absent when the turn
-  // produced no assistant message with a responseId (e.g. an aborted turn
-  // before any model call) — never guessed client-side (risk #5).
-  | { type: "turn_end"; lane?: Lane }
-  | { type: "queue_update"; steering: readonly string[]; followUp: readonly string[] }
-  // Replay a session's transcript (rebuilds the thread); and the sidebar list.
-  | { type: "history"; items: HistoryItem[] }
-  | { type: "sessions"; items: SessionListItem[]; activePath?: string }
-  // Sibling group for the LAST user message (edit-and-resend's `< i/n >`
-  // toggle, plan §5.2). Sent after every history rebuild and after a
-  // completed turn. `count` <= 1 means no toggle to show (the common case:
-  // no edits/regenerations yet). `entryId` is the currently-active sibling's
-  // id; `siblingIds` is the full ordered group so the client can resolve
-  // "previous/next" locally and send switch_sibling with a concrete id
-  // (never inferred/guessed — the server is the source of truth for order).
-  | { type: "siblings"; entryId?: string; index: number; count: number; siblingIds: string[] }
-  // Context-window usage indicator. tokens/percent are null right after a
-  // compaction until the next assistant reply (pi can't estimate yet).
-  | { type: "context"; tokens: number | null; contextWindow: number; percent: number | null }
-  // Effective vs. pending codingTools state (set_coding_tools above).
-  // `active` is what the CURRENT session actually built with (gates
-  // whether the approval card can ever fire); `pending` is the browser's
-  // most recent request, which only takes effect on the next
-  // new/opened/forked session — surfaced so the settings copy can say
-  // "will apply to your next new chat" instead of lying about immediacy.
-  | { type: "coding_tools"; active: boolean; pending: boolean }
-  // Full always-allow set for THIS machine (src/tool-approvals.ts), sent on
-  // `ready` and after any change, so the settings panel can list/forget
-  // durable per-tool approvals without a separate REST round-trip.
-  | { type: "tool_approvals"; alwaysAllow: string[] }
-  // App-aware assistant (plan §6.6): the navigate_app / spotlight_ui tools'
-  // side-channel notifications, mapped 1:1 from the tool call's params —
-  // never guessed or re-derived, so chat.ts's router/spotlight get exactly
-  // what the model asked for. `route` here is always a validated
-  // AppRouteId (isAppRouteId already ran server-side before this frame is
-  // ever sent — see navigate_app's execute). ui_spotlight's fields mirror
-  // SpotlightRequest 1:1 (assistant.ts resolves ref/label/selector/target
-  // client-side, in that precedence order, against its last snapshot).
-  | { type: "ui_navigate"; route: AppRouteId }
-  | { type: "ui_spotlight"; ref?: string; label?: string; selector?: string; target?: string; message?: string; route?: AppRouteId }
-  | { type: "error"; message: string };
-
-// ---- Session serialization (pure, unit-tested) -----------------------
-
-/** A finished tool call as rendered in replayed history. */
-export interface HistoryToolItem {
-  callId: string;
-  name: string;
-  args: unknown;
-  /** Tool result text, filled from the matching toolResult message. */
-  result: string;
-}
-
-/** One replayed transcript turn: a user or assistant message (assistant
- *  messages carry any tool calls they made, with results merged in). */
-export interface HistoryItem {
-  role: "user" | "assistant";
-  text: string;
-  /** Model reasoning/thinking, kept separate from the final answer. */
-  thinking?: string;
-  tools: HistoryToolItem[];
-  /** Session entry id (user items only). Lets the browser correlate "the
-   *  last user message" in the replayed thread with the `siblings` frame's
-   *  entryId, so it knows which DOM node gets the edit/sibling-toggle UI
-   *  (plan §5.2). Undefined for assistant items — they aren't fork targets. */
-  entryId?: string;
-}
-
-/** A row in the recent-chats sidebar, derived from pi's SessionInfo. */
-export interface SessionListItem {
-  path: string;
-  id: string;
-  title: string;
-  /** Last-modified epoch ms (newest first). */
-  modified: number;
-  messageCount: number;
-  forked: boolean;
-}
-
 /** Flatten a content value (string | content-parts) to its text. */
 function contentText(content: unknown): string {
   if (typeof content === "string") return content;
@@ -964,25 +676,6 @@ export function injectAdapter(
   return { ...payload, adapter: selected };
 }
 
-/** Per-request sampling overrides carried on the session and injected into the
- *  provider payload by the before_provider_request hook. Each field is either a
- *  user-set override or null/undefined ("use the server's mode-aware default").
- *  Field names match the server's ChatRequest wire names (not toOptions' camelCase)
- *  since these are injected directly into the outgoing chat-completions payload. */
-export interface SamplingOverrides {
-  temperature?: number | null;
-  top_p?: number | null;
-  top_k?: number | null;
-  min_p?: number | null;
-  xtc_probability?: number | null;
-  xtc_threshold?: number | null;
-  repetition_penalty?: number | null;
-  repetition_context_size?: number | null;
-  presence_penalty?: number | null;
-  frequency_penalty?: number | null;
-  seed?: number | null;
-}
-
 /** Every SamplingOverrides field name, in wire order — the single source of
  *  truth `injectSampling` and the `set_sampling` handler both walk, so a new
  *  field only needs adding here + the interface above. */
@@ -1158,7 +851,7 @@ function appAwareTextResult(text: string, details: unknown = {}): { content: [{
  *  null before the first one arrives); `notify` sends a ServerMessage
  *  (ui_navigate/ui_spotlight) to that same browser tab. Pure enough to
  *  unit test the tool -> frame mapping via a fake notify collecting sent
- *  frames (tests/pi-web.test.ts). */
+ *  frames (tests/research/pi-web.test.ts). */
 export function createAppAwareTools(
   getContext: () => AppUiContext | null,
   notify: (msg: ServerMessage) => void,
@@ -1267,7 +960,7 @@ class PiWebSession {
    *  "Per-message sampling scope" block comment above composeSampling for
    *  the full lifecycle. Wrapped in SamplingScopeState (not two loose
    *  fields) so the exact same applySetSampling/consumeForRequest pure
-   *  functions run here and in tests/pi-web.test.ts. */
+   *  functions run here and in tests/research/pi-web.test.ts. */
   private samplingScope: SamplingScopeState = initialSamplingScopeState();
   /** User's custom system-prompt text for this connection (set via
    *  set_system_prompt), null = none. Read by the before_agent_start hook