mlx-bun 0.0.5 → 0.0.6

package/src/memory/wikify.ts

@@ -0,0 +1,23 @@
+// mlx-bun memory — synthesis stage 5: WIKIFY (deterministic).  ⚠️ M1 STUB.
+//
+// Non-model editorial cleanup after synthesis: canonicalize [[wikilinks]]
+// (spaces↔underscores, case, resolve to existing stems), repair footnotes, and
+// run the deterministic editorial gate. Non-fatal — a failure here logs and
+// continues; it never corrupts an article. Mirrors lucien's wikify.ts +
+// normalize-wikilinks.ts + normalize-footnotes.ts.
+//
+// See docs/design/memory-system.md → "The nightly pipeline" (wikify/normalize).
+
+import type { SynthesisEvent } from "./pipeline";
+
+export interface WikifyResult {
+  articlesNormalized: number;
+  linksFixed: number;
+}
+
+/** Deterministic editorial gate + link/footnote normalize. STUB. */
+export async function wikify(_onEvent?: (e: SynthesisEvent) => void): Promise<WikifyResult> {
+  // TODO(M1): resolveWikilinkToStem across touched articles (vault.ts already
+  // has the resolver), footnote renumber/repair, {{stub}} marking.
+  throw new Error("memory.wikify is not implemented yet (M1)");
+}

package/src/mlx/metal-kernel.ts

@@ -67,6 +67,11 @@ export class MetalKernel {
   apply(inputs: MlxArray[], call: MetalKernelCall): MlxArray[] {
     if (this.#disposed) throw new Error("MetalKernel used after dispose");
     const cfg = C.mlx_fast_metal_kernel_config_new();
+    // bun:ffi ptr() does NOT retain its argument, so any JS buffer whose pointer
+    // crosses an FFI call below is anchored here until that call returns —
+    // otherwise GC could free it mid-call (the lifetime class behind the
+    // flash-CCE pin bug). Pure-additive: holds references, changes no logic.
+    const keepAlive: unknown[] = [];
     try {
       // Derive outputs + grid from the (current) input shapes when a function is
       // given — so the SAME kernel object answers correctly for any input shape
@@ -76,6 +81,7 @@ export class MetalKernel {
       const grid = typeof call.grid === "function" ? call.grid(inputs) : call.grid;
       for (const o of outputs) {
         const shape = new Int32Array(o.shape);
+        keepAlive.push(shape);
         if (C.mlx_fast_metal_kernel_config_add_output_arg(cfg, ptr(shape), BigInt(o.shape.length), o.dtype) !== 0)
           throw new Error(`metal_kernel add_output_arg failed: ${takeMlxError() ?? ""}`);
       }
@@ -85,19 +91,26 @@ export class MetalKernel {
         C.mlx_fast_metal_kernel_config_set_init_value(cfg, call.initValue);
       for (const [k, v] of Object.entries(call.templateInts ?? {})) {
         const kb = Buffer.from(k + "\0", "utf8");
+        keepAlive.push(kb);
         C.mlx_fast_metal_kernel_config_add_template_arg_int(cfg, ptr(kb), v);
       }
       for (const [k, v] of Object.entries(call.templateDtypes ?? {})) {
         const kb = Buffer.from(k + "\0", "utf8");
+        keepAlive.push(kb);
         C.mlx_fast_metal_kernel_config_add_template_arg_dtype(cfg, ptr(kb), v);
       }
 
       const handles = new BigUint64Array(inputs.map((a) => a.handle));
+      keepAlive.push(handles);
       const inVec = C.mlx_vector_array_new_data(ptr(handles), BigInt(inputs.length));
       const slot = new BigUint64Array([C.mlx_vector_array_new()]);
+      keepAlive.push(slot);
       const slotPtr = ptr(slot);
       const status = C.mlx_fast_metal_kernel_apply(slotPtr, this.#kernel, inVec, cfg, call.stream ?? gpuStream);
       const outVec = read.u64(slotPtr, 0);
+      // anchor: all ptr()'d JS buffers above are done being read by native code by
+      // this point; touching keepAlive keeps the JIT from freeing them earlier.
+      if (keepAlive.length < 0) throw new Error("unreachable");
       C.mlx_vector_array_free(inVec);
       if (status !== 0) {
         C.mlx_vector_array_free(outVec);

package/src/model/flash-attention.ts

@@ -29,10 +29,12 @@ function forwardTiles(headDim: number): [number, number] {
   throw new Error(`unsupported head_dim=${headDim}`);
 }
 
-/** uint32 [n] array (matches optiq's shape/causal input dtype). */
+/** uint32 [n] array (matches optiq's shape/causal input dtype). COPY into an
+ *  mlx-owned array (not the pinning fromView) — these tiny kernel-arg arrays are
+ *  disposed before the lazy kernel evaluates; pinning leaks the host buffer per
+ *  call (see flash-cce.ts u32 for the full diagnosis). */
 function u32(values: number[]): MlxArray {
-  const u = new Uint32Array(values);
-  return MlxArray.fromView(new Uint8Array(u.buffer.slice(0)), [values.length], Dtype.uint32);
+  return MlxArray.fromBytesCopy(new Uint8Array(new Uint32Array(values).buffer), [values.length], Dtype.uint32);
 }
 
 // Forward shader — verbatim from optiq flash_attention_metal.py _KERNEL_SOURCE.

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 {