mlx-bun 0.0.4 → 0.0.6

package/src/model/gemma4.ts

@@ -51,22 +51,48 @@ import {
   type SharedKv,
   type LoraWeights,
 } from "./gemma4-base";
-import { fusedGeglu, fusedGeluEnabled, fusedGeluSupported } from "./fused-geglu-kernel";
+import { fusedGeglu, fusedGegluDifferentiable, fusedGeluEnabled, fusedGeluSupported, fusedGeluTraining } from "./fused-geglu-kernel";
 import { Checkpoint } from "../mlx/checkpoint";
 import { flashAttention, getTrainingAttn, flashSupported } from "./flash-attention";
 
 /** Optional gradient-checkpointing context for the training forward. When set,
  *  forwardLayers wraps each layer in a Checkpoint so its interior activations
  *  are recomputed in the backward pass. `byLayer` maps a layer index to that
- *  layer's trainable LoRA weights (re-swapped inside the closure so they are
- *  explicit checkpoint inputs — required because the autograd primals are
- *  disposed before the recompute). `keepAlive` collects the per-layer
+ *  layer's trainable LoRA weights, partitioned into the attention sub-block
+ *  (q/k/v/o_proj) and the MLP sub-block (gate/up/down_proj + per-layer gate/
+ *  projection) — re-swapped inside the closure so they are explicit checkpoint
+ *  inputs (required because the autograd primals are disposed before the
+ *  recompute). When `splitMlp` is set, the layer is wrapped in TWO checkpoints
+ *  (attn sub-block + MLP sub-block) with the post-attn residual `hMid` as the
+ *  boundary, so the backward recompute holds `max(attn,MLP)+hMid` instead of the
+ *  whole layer (the intra-layer MLP split, lever 4). `keepAlive` collects the
  *  Checkpoint objects; the trainer disposes them after value_and_grad. */
+export interface LayerLoras {
+  attn: LoraWeights[];
+  mlp: LoraWeights[];
+}
 export interface GradCheckpointCtx {
-  byLayer: Map<number, LoraWeights[]>;
+  byLayer: Map<number, LayerLoras>;
+  splitMlp: boolean;
   keepAlive: Checkpoint[];
 }
 
+// Shared prompt-prefix plan (lever 7, e4b — docs/design/orpo-training.md). When
+// set, Attention ropes the concatenated [prompt(P); chosenResp(Rc); rejectedResp(Rr)]
+// sequence BLOCK-WISE — prompt at offset 0, EACH response reset to logical offset
+// P — instead of the uniform cache.offset, so ONE forward over the concat is
+// bit-exact with the two separate [prompt;resp] forwards (each response sees the
+// prompt's RoPE positions, not shifted-by-the-other-response). The matching
+// block-sparse + LOGICAL-window attention mask rides in via the prefix cache's
+// makeMask (src/train/prefix-shared.ts). Set around the single forward
+// (single-threaded), cleared after; null → the normal uniform-offset rope (every
+// other forward is untouched).
+export interface GemmaPrefixPlan { P: number; Rc: number; Rr: number }
+let _gemmaPrefixPlan: GemmaPrefixPlan | null = null;
+export function setGemmaPrefixPlan(p: GemmaPrefixPlan | null): void {
+  _gemmaPrefixPlan = p;
+}
+
 class Attention {
   readonly isSliding: boolean;
   readonly windowSize: number;
@@ -161,6 +187,32 @@ class Attention {
       : ops.ropeDynamic(x, this.headDim, this.ropeBase, offset, this.ropeFreqs);
   }
 
+  /** Block-wise RoPE for the prefix-shared concat [prompt; chosen; rejected] along
+   *  the sequence axis (axis 2 of [B,H,T,D]): prompt at offset 0, EACH response
+   *  reset to logical offset P. RoPE is per-token, so roping each contiguous block
+   *  at its scalar offset and concatenating == roping with per-token position-ids
+   *  [0..P-1, P..P+Rc-1, P..P+Rr-1]. Uses this.rope per block so the layer's
+   *  proportional/ropeFreqs + sliding-vs-full geometry applies identically to the
+   *  two-forward path. Caller disposes the input. */
+  ropeBlocks(x: MlxArray, plan: GemmaPrefixPlan): MlxArray {
+    const { P, Rc, Rr } = plan;
+    const [B, H, , D] = x.shape as [number, number, number, number];
+    const blocks = [
+      { start: 0, len: P, off: 0 },
+      { start: P, len: Rc, off: P },
+      { start: P + Rc, len: Rr, off: P },
+    ].filter((b) => b.len > 0);
+    const parts = blocks.map((b) => {
+      const sl = x.slice([0, 0, b.start, 0], [B, H, b.start + b.len, D]);
+      const r = this.rope(sl, b.off);
+      sl.dispose();
+      return r;
+    });
+    const out = ops.concatAxis(parts, 2);
+    for (const p of parts) p.dispose();
+    return out;
+  }
+
   /** Returns the attention output plus the fetched KV (for KV-shared
    *  sharer layers and the speculative drafter). The SharedKv arrays are
    *  owned by the caller (forwardLayers) and disposed after the pass. */
@@ -203,7 +255,9 @@ class Attention {
       const kNormed = this.kNorm!.forward(k);
       const kT = ops.transposeAxes(kNormed, [0, 2, 1, 3]);
       kNormed.dispose();
-      const kRoped = this.rope(kT, offsetArr ?? offset);
+      const kRoped = _gemmaPrefixPlan
+        ? this.ropeBlocks(kT, _gemmaPrefixPlan)
+        : this.rope(kT, offsetArr ?? offset);
       kT.dispose();
 
       const vNormed = this.vNorm!.forward(v);
@@ -230,7 +284,9 @@ class Attention {
     }
 
     q = disposing(q, ops.transposeAxes(q, [0, 2, 1, 3]));
-    q = disposing(q, this.rope(q, shared.offsetArr ?? shared.offset));
+    q = disposing(q, _gemmaPrefixPlan
+      ? this.ropeBlocks(q, _gemmaPrefixPlan)
+      : this.rope(q, shared.offsetArr ?? shared.offset));
 
     let attn: MlxArray;
     const ta = getTrainingAttn();
@@ -273,9 +329,12 @@ class MLP {
     const u = this.up.forward(x);
     // Fused GeGLU perf kernel (opt-in, not in a compiled trace): one pass
     // instead of ~9 element-wise kernels. Off → the bit-exact spelled path.
+    // In training mode the differentiable wrapper (kernel forward + hand-derived
+    // vjp) is used so autograd can flow AND the backward recomputes the gelu
+    // from the primal instead of retaining the spelled-out intermediates.
     let h: MlxArray;
     if (fusedGeluEnabled() && !isCompiledTrace() && fusedGeluSupported(g)) {
-      h = fusedGeglu(g, u);
+      h = fusedGeluTraining() ? fusedGegluDifferentiable(g, u) : fusedGeglu(g, u);
       g.dispose();
       u.dispose();
     } else {
@@ -500,6 +559,23 @@ export class DecoderLayer {
   forward(
     x: MlxArray, mask: Mask, cache: Cache | null,
     sharedIn: SharedKv | null, perLayerInput: MlxArray | null,
+  ): { h: MlxArray; shared: SharedKv } {
+    // Composed from the two sub-blocks (the intra-layer min-cut is the post-attn
+    // residual `hMid` [1,T,hidden]). Behaviour is identical to the spelled-out
+    // body; the split lets the checkpoint path rematerialize attn and MLP
+    // separately (peak attn+MLP → max(attn,MLP)+hMid). `forwardMlp` borrows hMid.
+    const { h: hMid, shared } = this.forwardAttn(x, mask, cache, sharedIn);
+    const h = this.forwardMlp(hMid, perLayerInput);
+    hMid.dispose();
+    return { h, shared };
+  }
+
+  /** Attention sub-block: inputNorm → attn → postAttnNorm → residual add.
+   *  Returns the post-attn residual stream `h` [1,T,hidden] (the intra-layer
+   *  boundary) and the attention's fetched/produced K/V (`shared`). `x` is
+   *  borrowed (never disposed). The e4b donor-KV threading lives entirely here. */
+  forwardAttn(
+    x: MlxArray, mask: Mask, cache: Cache | null, sharedIn: SharedKv | null,
   ): { h: MlxArray; shared: SharedKv } {
     let h = this.inputNorm.forward(x);
     const { out, shared } = this.attn.forward(h, mask, cache, sharedIn);
@@ -507,18 +583,24 @@ export class DecoderLayer {
     h = out;
     h = disposing(h, this.postAttnNorm.forward(h));
     h = disposing(h, ops.add(x, h));
+    return { h, shared };
+  }
 
-    const residual = h;
+  /** MLP sub-block: FFN (dense or MoE) + postFfNorm + residual add, then the
+   *  e4b per-layer-input gate and layer scalar. `hMid` (the attn sub-block's
+   *  output) is BORROWED — the caller owns/disposes it (it is the checkpoint
+   *  boundary leaf), so this never disposes it. */
+  forwardMlp(hMid: MlxArray, perLayerInput: MlxArray | null): MlxArray {
     let f: MlxArray;
     if (this.router && this.experts) {
-      // MoE: dense MLP and routed experts as parallel branches off h
+      // MoE: dense MLP and routed experts as parallel branches off hMid
       // (reference DecoderLayer, enable_moe path)
-      let h1 = this.preFfNorm.forward(h);
+      let h1 = this.preFfNorm.forward(hMid);
       h1 = disposing(h1, this.mlp.forward(h1));
       h1 = disposing(h1, this.postFfNorm1!.forward(h1));
 
-      const { indices, weights: topKWeights } = this.router.forward(h);
-      let h2 = this.preFfNorm2!.forward(h);
+      const { indices, weights: topKWeights } = this.router.forward(hMid);
+      let h2 = this.preFfNorm2!.forward(hMid);
       h2 = disposing(h2, this.experts.forward(h2, indices, topKWeights));
       h2 = disposing(h2, this.postFfNorm2!.forward(h2));
       indices.dispose();
@@ -528,20 +610,25 @@ export class DecoderLayer {
       h1.dispose();
       h2.dispose();
     } else {
-      f = this.preFfNorm.forward(h);
+      f = this.preFfNorm.forward(hMid);
       f = disposing(f, this.mlp.forward(f));
     }
     f = disposing(f, this.postFfNorm.forward(f));
-    h = ops.add(residual, f);
-    residual.dispose();
+    let h = ops.add(hMid, f); // hMid borrowed — NOT disposed here
     f.dispose();
 
-    // per-layer input gating (reference gemma4_text DecoderLayer)
+    // per-layer input gating (reference gemma4_text DecoderLayer). gelu(gate)·pli
+    // is the second GeGLU site — fused (differentiable) on the training path, the
+    // bit-exact spelled-out gelu+mul on the inference path (unchanged).
     if (this.perLayerGate && this.perLayerProjection && this.postPerLayerNorm && perLayerInput) {
       const res2 = h;
       let gate = this.perLayerGate.forward(h);
-      gate = disposing(gate, ops.geluApprox(gate));
-      gate = disposing(gate, ops.mul(gate, perLayerInput));
+      if (fusedGeluTraining() && fusedGeluEnabled() && !isCompiledTrace() && fusedGeluSupported(gate)) {
+        gate = disposing(gate, fusedGegluDifferentiable(gate, perLayerInput));
+      } else {
+        gate = disposing(gate, ops.geluApprox(gate));
+        gate = disposing(gate, ops.mul(gate, perLayerInput));
+      }
       gate = disposing(gate, this.perLayerProjection.forward(gate));
       gate = disposing(gate, this.postPerLayerNorm.forward(gate));
       h = ops.add(res2, gate);
@@ -550,7 +637,7 @@ export class DecoderLayer {
     }
 
     if (this.layerScalar) h = disposing(h, ops.mul(h, this.layerScalar));
-    return { h, shared };
+    return h;
   }
 }
 
@@ -927,7 +1014,19 @@ export class Gemma4Model {
     layerCache: Cache | null, sharedIn: SharedKv | null, pls: MlxArray | null,
   ): { h: MlxArray; shared: SharedKv | null } {
     const ctx = this.gradCkpt!;
-    const lws = ctx.byLayer.get(i) ?? [];
+    const ll = ctx.byLayer.get(i) ?? { attn: [], mlp: [] };
+    return ctx.splitMlp
+      ? this.runSplitCheckpointedLayer(layer, h, mask, layerCache, sharedIn, pls, ll, ctx)
+      : this.runSingleCheckpointedLayer(layer, h, mask, layerCache, sharedIn, pls, ll, ctx);
+  }
+
+  /** One Checkpoint over the whole layer (the original path). */
+  private runSingleCheckpointedLayer(
+    layer: DecoderLayer, h: MlxArray, mask: Mask,
+    layerCache: Cache | null, sharedIn: SharedKv | null, pls: MlxArray | null,
+    ll: LayerLoras, ctx: GradCheckpointCtx,
+  ): { h: MlxArray; shared: SharedKv | null } {
+    const lws = [...ll.attn, ...ll.mlp];
     const isSharer = layerCache === null;
     const sIn = isSharer && sharedIn && sharedIn.kind === "plain" ? sharedIn : null;
     const hasMaskArr = mask.arr !== null;
@@ -970,6 +1069,82 @@ export class Gemma4Model {
     return { h: outs[0]!, shared };
   }
 
+  /** Intra-layer MLP split (lever 4): TWO checkpoints — the attention sub-block
+   *  (forwardAttn: inputNorm → attn → residual add, threading the donor K/V) and
+   *  the MLP sub-block (forwardMlp: FFN + per-layer gate) — with the post-attn
+   *  residual `hMid` as the boundary. The backward recompute holds only one
+   *  sub-block's activations at a time, dropping per-layer peak from `attn+MLP`
+   *  to `max(attn,MLP)+hMid`. Attn LoRA / pls / donor-KV go to the attn ck; MLP
+   *  LoRA to the MLP ck. */
+  private runSplitCheckpointedLayer(
+    layer: DecoderLayer, h: MlxArray, mask: Mask,
+    layerCache: Cache | null, sharedIn: SharedKv | null, pls: MlxArray | null,
+    ll: LayerLoras, ctx: GradCheckpointCtx,
+  ): { h: MlxArray; shared: SharedKv | null } {
+    const isSharer = layerCache === null;
+    const sIn = isSharer && sharedIn && sharedIn.kind === "plain" ? sharedIn : null;
+    const hasMaskArr = mask.arr !== null;
+    const hasPls = pls !== null;
+    const maskMode = mask.mode;
+    const attnLws = ll.attn;
+    const mlpLws = ll.mlp;
+
+    // --- attn sub-block checkpoint: [h, maskArr?, sharedInK?, sharedInV?, attnLora*]
+    const attnInputs: MlxArray[] = [h];
+    if (hasMaskArr) attnInputs.push(mask.arr!);
+    if (sIn) attnInputs.push(sIn.keys, sIn.values);
+    for (const lw of attnLws) attnInputs.push(lw.a, lw.b);
+
+    const attnCk = new Checkpoint((ins) => {
+      let k = 0;
+      const hIn = ins[k++]!;
+      const mArr = hasMaskArr ? ins[k++]! : null;
+      const sharedInRebuilt: SharedKv | null = sIn
+        ? { kind: "plain", keys: ins[k++]!, values: ins[k++]!, offset: 0 }
+        : null;
+      const saved = attnLws.map((lw) => [lw.a, lw.b] as [MlxArray, MlxArray]);
+      for (const lw of attnLws) { lw.a = ins[k++]!; lw.b = ins[k++]!; }
+      try {
+        const out = layer.forwardAttn(hIn, { mode: maskMode, arr: mArr }, layerCache, sharedInRebuilt);
+        if (!isSharer && out.shared && out.shared.kind === "plain")
+          return [out.h, out.shared.keys, out.shared.values];
+        return [out.h];
+      } finally {
+        attnLws.forEach((lw, j) => { lw.a = saved[j]![0]; lw.b = saved[j]![1]; });
+      }
+    });
+    ctx.keepAlive.push(attnCk);
+    const attnOuts = attnCk.apply(attnInputs);
+    const hMid = attnOuts[0]!;
+    const shared: SharedKv | null =
+      !isSharer && attnOuts.length === 3
+        ? { kind: "plain", keys: attnOuts[1]!, values: attnOuts[2]!, offset: 0 }
+        : null;
+
+    // --- MLP sub-block checkpoint: [hMid, pls?, mlpLora*] -> [hOut]
+    const mlpInputs: MlxArray[] = [hMid];
+    if (hasPls) mlpInputs.push(pls!);
+    for (const lw of mlpLws) mlpInputs.push(lw.a, lw.b);
+
+    const mlpCk = new Checkpoint((ins) => {
+      let k = 0;
+      const hMidIn = ins[k++]!;
+      const plsIn = hasPls ? ins[k++]! : null;
+      const saved = mlpLws.map((lw) => [lw.a, lw.b] as [MlxArray, MlxArray]);
+      for (const lw of mlpLws) { lw.a = ins[k++]!; lw.b = ins[k++]!; }
+      try {
+        return [layer.forwardMlp(hMidIn, plsIn)];
+      } finally {
+        mlpLws.forEach((lw, j) => { lw.a = saved[j]![0]; lw.b = saved[j]![1]; });
+      }
+    });
+    ctx.keepAlive.push(mlpCk);
+    const mlpOuts = mlpCk.apply(mlpInputs);
+    hMid.dispose(); // boundary leaf consumed by the MLP ck (mlx retains it in-graph)
+
+    return { h: mlpOuts[0]!, shared };
+  }
+
   /** hidden [1, L, hidden] → softcapped logits [1, L, vocab]. */
   logitsFromHidden(h: MlxArray): MlxArray {
     let logits = this.embed.asLinear(h);

package/src/model/minicpm5.ts

@@ -21,6 +21,45 @@ import {
 } from "./gemma4-base";
 import { flashAttention, getTrainingAttn, flashSupported } from "./flash-attention";
 
+// Shared prompt-prefix plan (lever 7). When set, the attention ropes the
+// concatenated [prompt(P); chosenResp(Rc); rejectedResp(Rr)] sequence
+// BLOCK-WISE — prompt at offset 0, EACH response reset to offset P — instead of
+// the uniform cache.offset. This is what lets one forward over the concat be
+// bit-exact with the two separate [prompt;resp] forwards (each response sees the
+// prompt's RoPE positions, not shifted-by-the-other-response). Set by
+// orpoLossPrefixShared around the forward (single-threaded), cleared after; null
+// → the normal uniform-offset rope (every other forward is untouched). The
+// matching block-sparse attention mask rides in via PrefixSharedCache.makeMask.
+export interface PrefixPlan { P: number; Rc: number; Rr: number }
+let _prefixPlan: PrefixPlan | null = null;
+export function setMiniCpmPrefixPlan(p: PrefixPlan | null): void {
+  _prefixPlan = p;
+}
+
+/** Block-wise RoPE for the prefix-shared concat [prompt; chosen; rejected] along
+ *  the sequence axis (axis 2 of [B,H,T,D]): prompt rotated at offset 0, each
+ *  response at offset P (reset). RoPE is per-token, so roping each contiguous
+ *  block at its scalar offset and concatenating == roping with per-token
+ *  position-ids [0..P-1, P..P+Rc-1, P..P+Rr-1]. Caller disposes the input. */
+function ropeBlocks(x: MlxArray, dims: number, base: number, plan: PrefixPlan): MlxArray {
+  const { P, Rc, Rr } = plan;
+  const [B, H, , D] = x.shape as [number, number, number, number];
+  const blocks: { start: number; len: number; off: number }[] = [
+    { start: 0, len: P, off: 0 },
+    { start: P, len: Rc, off: P },
+    { start: P + Rc, len: Rr, off: P },
+  ].filter((b) => b.len > 0);
+  const parts = blocks.map((b) => {
+    const sl = x.slice([0, 0, b.start, 0], [B, H, b.start + b.len, D]);
+    const r = ops.rope(sl, dims, base, b.off, null);
+    sl.dispose();
+    return r;
+  });
+  const out = ops.concatAxis(parts, 2);
+  for (const p of parts) p.dispose();
+  return out;
+}
+
 class LlamaAttention {
   readonly qProj: QuantizedLinear;
   readonly kProj: QuantizedLinear;
@@ -65,9 +104,11 @@ class LlamaAttention {
     // runs before updateAndFetch, so K and Q share the pre-write offset.
     const offsetArr = cache.ropeOffsetArr;
     const ropeStep = (x: MlxArray): MlxArray =>
-      offsetArr
-        ? ops.ropeDynamic(x, this.headDim, this.ropeBase, offsetArr, null)
-        : ops.rope(x, this.headDim, this.ropeBase, cache.offset, null);
+      _prefixPlan
+        ? ropeBlocks(x, this.headDim, this.ropeBase, _prefixPlan)
+        : offsetArr
+          ? ops.ropeDynamic(x, this.headDim, this.ropeBase, offsetArr, null)
+          : ops.rope(x, this.headDim, this.ropeBase, cache.offset, null);
     q = disposing(q, ropeStep(q));
     k = disposing(k, ropeStep(k));
     let attn: MlxArray;

package/src/pi-provider.ts

@@ -40,6 +40,13 @@ export interface PiProviderOptions {
    *  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;
 }
 
 export interface PiProviderWiring {
@@ -69,19 +76,28 @@ export function buildPiProvider(baseUrl: string, opts: PiProviderOptions = {}):
         name: opts.name ?? "mlx-bun (local)",
         api: PI_API,
         reasoning: opts.reasoning ?? false,
-        // Our local thinking models (Qwen3.5/MiniCPM5) are BINARY: the chat
-        // template's enable_thinking is a boolean. "qwen-chat-template" makes
-        // Pi send chat_template_kwargs.enable_thinking; pi-ai computes that
-        // boolean from whether the selected level is off or not, so no level
-        // map is needed here (null in thinkingLevelMap means "unsupported").
-        ...(opts.reasoning
-          ? { compat: { thinkingFormat: "qwen-chat-template" as const } }
-          : {}),
-        // Declare image capability so pi will carry image content to /v1.
-        // Whether the *loaded* model can actually see images is gated at the
-        // server (non-vision models 400) and surfaced to the web UI via the
-        // `ready` frame's `vision` flag.
-        input: ["text", "image"],
+        // 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,

package/src/pi-session.ts

@@ -0,0 +1,81 @@
+// Shared pi front-door assembly for mlx-bun's embedded web + terminal agents.
+//
+// This is the single place that decides which custom tools and bundled skills
+// are visible to pi. Keeping that policy here prevents the browser and TUI from
+// drifting — especially for personal memory, whose prompt surface must be
+// relevant-but-not-dominant.
+
+import type { ToolDefinition } from "@earendil-works/pi-coding-agent";
+import { createWebTools, WEB_TOOL_NAMES } from "./web-tools";
+import { materializeBundledSkillPaths } from "./web/skills";
+import {
+  createMemoryTools,
+  MEMORY_TOOL_NAMES,
+  memoryIndexHint,
+  isMemoryEnabled,
+  createReferenceTools,
+  REFERENCE_TOOL_NAMES,
+} from "./memory/tools";
+
+/** Built-in pi coding tools plus mlx-bun's non-mutating web tools. */
+export const PI_BASE_TOOL_NAMES = ["read", "bash", "edit", "write", "grep", "find", "ls", ...WEB_TOOL_NAMES] as const;
+
+export interface PiAgentSurfaceOptions {
+  /** Force memory on/off for tests. Default: auto-detect vault existence. */
+  memory?: "auto" | "on" | "off";
+  /** Include web_search/web_fetch/weather. Default true. */
+  webTools?: boolean;
+  /** Include read/bash/edit/write/grep/find/ls. Default true. */
+  codingTools?: boolean;
+}
+
+export interface PiAgentSurface {
+  memoryEnabled: boolean;
+  /** Tool allowlist to pass to createAgentSession/createAgentSessionFromServices. */
+  tools: string[];
+  /** Custom tools to pass to pi. Built-ins are selected by name in `tools`. */
+  customTools: ToolDefinition[];
+  /** Skill paths for DefaultResourceLoader.additionalSkillPaths. */
+  skillPaths: string[];
+  /** Short prompt addendum. Empty unless memory is enabled. */
+  memoryHint: string;
+}
+
+async function resolveMemoryEnabled(mode: PiAgentSurfaceOptions["memory"]): Promise<boolean> {
+  if (mode === "on") return true;
+  if (mode === "off") return false;
+  return isMemoryEnabled();
+}
+
+/**
+ * Build the shared tool + skill surface for both pi front doors.
+ *
+ * Memory policy: if the user enabled memory (vault exists), memory is available.
+ * It is still a scoped retrieval capability, not a default first step for every
+ * prompt. If memory is absent, no memory skill/tool schemas are exposed at all.
+ */
+export async function buildPiAgentSurface(opts: PiAgentSurfaceOptions = {}): Promise<PiAgentSurface> {
+  const memoryEnabled = await resolveMemoryEnabled(opts.memory ?? "auto");
+  const includeWeb = opts.webTools ?? true;
+  const includeCoding = opts.codingTools ?? true;
+
+  const tools: string[] = [];
+  if (includeCoding) tools.push("read", "bash", "edit", "write", "grep", "find", "ls");
+  const customTools: ToolDefinition[] = [];
+  if (includeWeb) {
+    tools.push(...WEB_TOOL_NAMES);
+    customTools.push(...createWebTools());
+  }
+  if (memoryEnabled) {
+    tools.push(...MEMORY_TOOL_NAMES, ...REFERENCE_TOOL_NAMES);
+    customTools.push(...createMemoryTools(), ...createReferenceTools());
+  }
+
+  return {
+    memoryEnabled,
+    tools,
+    customTools,
+    skillPaths: materializeBundledSkillPaths({ includeMemory: memoryEnabled }),
+    memoryHint: memoryEnabled ? await memoryIndexHint() : "",
+  };
+}

package/src/pi-terminal.ts

@@ -35,8 +35,7 @@ import {
   runRpcMode,
   type CreateAgentSessionRuntimeFactory,
 } from "@earendil-works/pi-coding-agent";
-import { createWebTools, WEB_TOOL_NAMES } from "./web-tools";
-import { materializeBundledSkills } from "./web/skills";
+import { buildPiAgentSurface } from "./pi-session";
 import { buildPiProvider, DEFAULT_CONTEXT_WINDOW } from "./pi-provider";
 
 /** How the embedded agent runs. Mirrors pi's CLI appMode. */
@@ -49,8 +48,6 @@ export type PiTerminalMode = "interactive" | "print" | "rpc";
  * sees it. Mutating tools (bash/edit/write) surface pi's own TUI approval
  * prompt at call time — we don't add a gate.
  */
-const ALL_TOOLS = ["read", "bash", "edit", "write", "grep", "find", "ls", ...WEB_TOOL_NAMES];
-
 /**
  * System prompt for the embedded terminal coding agent.
  *
@@ -70,20 +67,11 @@ export function buildTerminalSystemPrompt(modelLabel?: string): string {
   const model = modelLabel ? `the local model \`${modelLabel}\`` : "a local model";
   return `You are mlx-bun's terminal coding agent — a capable, hands-on pair programmer working alongside the user in their terminal. You run on ${model} entirely on the user's own Apple-silicon Mac via mlx-bun, so the whole session stays local and private; nothing leaves the machine except the explicit web tools below.
 
-You operate inside the user's current project directory and have a real toolset — use it. Prefer doing the work over describing it: open the files, search the tree, run the command, make the edit. Investigate before you answer, and follow through to a working result.
-
-Tools used freely (no confirmation needed):
-- read — open and read any file
-- ls, find, grep — list directories and search the tree by name or content
-- web_search — search the web for current, real-world information (docs, APIs, errors, anything that changes or that you don't already know)
-- web_fetch — fetch a URL and read the page or data behind it (a natural follow-up to web_search, or when the user gives you a link)
-- weather — current conditions and a short forecast for any place
+You operate inside the user's current project directory. Available capabilities when relevant: read-only file inspection (read, ls, find, grep), web search/fetch (web_search, web_fetch), weather, and user-approved local actions (bash, edit, write).
 
-Tools that ask the user for approval first (they stay in control — pi shows the prompt, you don't have to ask in prose):
-- bash — run shell commands (build, test, git, scripts)
-- edit, write — modify existing files or create new ones
+Use tools in service of helping the user, not as a performance. For ordinary social conversation, answer naturally yourself. For real-world facts, current events, schedules, releases, local information, recommendations, restaurants, movies, prices, travel, docs, or anything likely to have changed, look it up with web_search/web_fetch or weather instead of relying on memory or guessing. For coding tasks, inspect the files, investigate before you answer, and follow through to a working result. When you need to change code or run something, propose the concrete edit or command and let the approval prompt handle consent — don't ask permission in prose first.
 
-When you need to change code or run something, just propose the concrete edit or command and let the approval prompt handle consent — don't narrate "I would run…" or ask permission in text. When you hit something you don't know — especially current or version-specific facts — look it up with web_search and web_fetch instead of guessing. If a request is genuinely ambiguous, ask one short clarifying question; otherwise make a sensible choice and proceed.
+Environment metadata such as current date, working directory, model name, and available tools may be provided for context; do not report or summarize that metadata unless the user asks for it or it is directly relevant. Do not narrate tool policies or internal workflows.
 
 Be honest about what you did and what you found, including mistakes, dead ends, and anything you couldn't do. Keep responses tight and skimmable: format with Markdown, show file paths and commands plainly, and cite links when you used the web.`;
 }
@@ -121,6 +109,7 @@ export async function runEmbeddedPi(opts: RunEmbeddedPiOptions): Promise<number>
   const cwd = process.cwd();
   // Own session/settings dir, isolated from any standalone pi install.
   const agentDir = join(homedir(), ".mlx-bun", "pi");
+  const sessionDir = join(agentDir, "sessions", `--${cwd.replace(/^[/\\]/, "").replace(/[/\\:]/g, "-")}--`);
 
   const { authStorage, modelRegistry, model } = buildPiProvider(opts.baseUrl, {
     contextWindow: opts.contextWindow ?? DEFAULT_CONTEXT_WINDOW,
@@ -128,10 +117,10 @@ export async function runEmbeddedPi(opts: RunEmbeddedPiOptions): Promise<number>
     reasoning: opts.reasoning ?? false,
   });
 
-  // Curated skills (web-research, …) written to ~/.mlx-bun/skills and loaded
-  // on top of an otherwise-isolated session. See src/web/skills.ts.
-  const skillsRoot = materializeBundledSkills();
-  const systemPrompt = buildTerminalSystemPrompt(opts.modelLabel);
+  // Shared pi tool/skill surface. Memory is included only when the user has
+  // enabled a vault; if present it remains scoped to user-specific continuity.
+  const surface = await buildPiAgentSurface();
+  const systemPrompt = buildTerminalSystemPrompt(opts.modelLabel) + surface.memoryHint;
 
   // The runtime factory: same shape as pi's own CLI (dist/main.js) and the
   // shipped examples/sdk/13-session-runtime.ts. The same factory is reused by
@@ -154,7 +143,7 @@ export async function runEmbeddedPi(opts: RunEmbeddedPiOptions): Promise<number>
         // AGENTS.md / CLAUDE.md context files (unlike the web chat).
         noContextFiles: false,
         // Our curated skills on top (noSkills above only disables discovery).
-        additionalSkillPaths: [skillsRoot],
+        additionalSkillPaths: surface.skillPaths,
         // Replace pi's default coding-agent prompt with our mlx-bun persona.
         systemPrompt,
       },
@@ -165,15 +154,15 @@ export async function runEmbeddedPi(opts: RunEmbeddedPiOptions): Promise<number>
         sessionManager,
         sessionStartEvent,
         model,
-        tools: ALL_TOOLS,
-        customTools: createWebTools(),
+        tools: surface.tools,
+        customTools: surface.customTools,
       })),
       services,
       diagnostics: services.diagnostics,
     };
   };
 
-  const sessionManager = SessionManager.create(cwd);
+  const sessionManager = SessionManager.create(cwd, sessionDir);
   const runtime = await createAgentSessionRuntime(createRuntime, { cwd, agentDir, sessionManager });
 
   try {