mlx-bun 0.0.10 → 0.0.12

package/src/memory/rest.ts

@@ -0,0 +1,245 @@
+// Memory REST wrappers — thin loopback JSON routes over src/memory/vault.ts
+// for the web chat's Memory panel (docs/design/web-chat-redesign.md §5.5,
+// §9 Phase 2). These are HTTP-only conveniences for the frontend: the
+// agent-tool surface (src/memory/tools.ts) stays read-only and unchanged,
+// and every route here is itself read-only except /init, which delegates to
+// the same `setupVault` the CLI's `mlx-bun memory init` calls (never
+// duplicated logic).
+//
+// Every handler returns a Response and never throws — a missing vault or a
+// bad param degrades to `{ ok: false, error }` (never a 500), matching the
+// rest of server.ts's /api/* conventions. Route dispatch (matching
+// url.pathname + method) lives in server.ts; this file is the pure handler
+// bodies so they're unit-testable without booting a model context (see
+// tests/memory-rest.test.ts).
+
+import { realpath } from "node:fs/promises";
+import { tmpdir } from "node:os";
+import { dirname, resolve, sep } from "node:path";
+import {
+  vaultRoot, vaultStatus, listArticles, listReferenceDocs, readArticle,
+  searchArticles, getArticleLinks, resolveArticleRelPath, articleHistory, articleDiff,
+  setupVault, type VaultStatus,
+} from "./vault";
+import { parseInfobox, parseLead, parseSeriesBanner, articleStructure } from "./article";
+
+function jsonOk<T extends object>(body: T, init?: ResponseInit): Response {
+  return Response.json({ ok: true, ...body }, init);
+}
+
+function jsonErr(error: string, status = 400): Response {
+  return Response.json({ ok: false, error }, { status });
+}
+
+/** `{ ok:false, enabled:false }` shape shared by every route when no vault
+ *  exists yet — never a throw, never a 404: the frontend uses this to
+ *  render the first-run consent card instead of an error toast. */
+function noVault(root: string): Response {
+  return Response.json({ ok: false, enabled: false, error: "no memory vault yet", root });
+}
+
+// ---- GET /api/memory/status -------------------------------------------
+
+export async function handleMemoryStatus(): Promise<Response> {
+  const st = await vaultStatus();
+  if (!st.exists) return noVault(st.root);
+  return jsonOk({ enabled: true, status: st satisfies VaultStatus });
+}
+
+// ---- GET /api/memory/list ----------------------------------------------
+
+export async function handleMemoryList(): Promise<Response> {
+  const root = vaultRoot();
+  const st = await vaultStatus(root);
+  if (!st.exists) return noVault(root);
+  const [articles, reference] = await Promise.all([listArticles(root), listReferenceDocs(root)]);
+  return jsonOk({ articles, reference });
+}
+
+// ---- GET /api/memory/search?q=&scope= -----------------------------------
+
+export async function handleMemorySearch(url: URL): Promise<Response> {
+  const root = vaultRoot();
+  const st = await vaultStatus(root);
+  if (!st.exists) return noVault(root);
+  const q = (url.searchParams.get("q") ?? "").trim();
+  if (!q) return jsonErr("q is required");
+  const scopeParam = url.searchParams.get("scope");
+  const scope = scopeParam === "articles" || scopeParam === "reference" ? scopeParam : "all";
+  const limitParam = Number(url.searchParams.get("limit") ?? "50");
+  const limit = Number.isFinite(limitParam) && limitParam > 0 ? Math.min(limitParam, 200) : 50;
+  try {
+    const { summaries, hits } = await searchArticles(root, q, { limit, scope });
+    return jsonOk({ summaries, hits });
+  } catch (e) {
+    return jsonErr((e as Error).message);
+  }
+}
+
+// ---- GET /api/memory/article?name= --------------------------------------
+//
+// Rendered source + metadata: the article's raw Markdown plus the
+// deterministic structure the panel needs to render it without re-parsing
+// client-side (infobox, lead, series banner, section skeleton).
+
+export async function handleMemoryArticle(url: URL): Promise<Response> {
+  const root = vaultRoot();
+  const st = await vaultStatus(root);
+  if (!st.exists) return noVault(root);
+  const name = (url.searchParams.get("name") ?? "").trim();
+  if (!name) return jsonErr("name is required");
+  try {
+    const { path, content } = await readArticle(root, name);
+    return jsonOk({
+      name,
+      path,
+      content,
+      infobox: parseInfobox(content),
+      lead: parseLead(content),
+      series: parseSeriesBanner(content),
+      structure: articleStructure(content),
+    });
+  } catch (e) {
+    return jsonErr((e as Error).message, 404);
+  }
+}
+
+// ---- GET /api/memory/links?name= -----------------------------------------
+
+export async function handleMemoryLinks(url: URL): Promise<Response> {
+  const root = vaultRoot();
+  const st = await vaultStatus(root);
+  if (!st.exists) return noVault(root);
+  const name = (url.searchParams.get("name") ?? "").trim();
+  if (!name) return jsonErr("name is required");
+  try {
+    const links = await getArticleLinks(root, name);
+    return jsonOk({ name, ...links });
+  } catch (e) {
+    return jsonErr((e as Error).message, 404);
+  }
+}
+
+// ---- GET /api/memory/history?name= ---------------------------------------
+
+export async function handleMemoryHistory(url: URL): Promise<Response> {
+  const root = vaultRoot();
+  const st = await vaultStatus(root);
+  if (!st.exists) return noVault(root);
+  const name = (url.searchParams.get("name") ?? "").trim();
+  if (!name) return jsonErr("name is required");
+  const limitParam = Number(url.searchParams.get("limit") ?? "50");
+  const limit = Number.isFinite(limitParam) && limitParam > 0 ? Math.min(limitParam, 200) : 50;
+  try {
+    const relPath = await resolveArticleRelPath(root, name); // throws if not found — validates `name`
+    const entries = await articleHistory(root, relPath, limit);
+    return jsonOk({ name, isGitRepo: st.isGitRepo, entries });
+  } catch (e) {
+    return jsonErr((e as Error).message, 404);
+  }
+}
+
+// ---- GET /api/memory/diff?name=&rev= -------------------------------------
+
+/** Full 40-hex or an abbreviated (≥4 hex) commit hash — never a ref
+ *  expression (`HEAD~1`, `main^`, `refs/...`), which could otherwise be
+ *  abused to walk outside the intended commit or (in principle) smuggle
+ *  option-like text into the argv. `rev` is passed to Bun.spawn as its own
+ *  argv element (never through a shell), but this check is the defense
+ *  against a caller trying to pass anything other than a bare commit id. */
+const REV_RE = /^[0-9a-f]{4,40}$/;
+
+export async function handleMemoryDiff(url: URL): Promise<Response> {
+  const root = vaultRoot();
+  const st = await vaultStatus(root);
+  if (!st.exists) return noVault(root);
+  const name = (url.searchParams.get("name") ?? "").trim();
+  const rev = (url.searchParams.get("rev") ?? "").trim();
+  if (!name) return jsonErr("name is required");
+  if (!rev) return jsonErr("rev is required");
+  if (!REV_RE.test(rev)) return jsonErr("rev must be a commit hash (4-40 hex chars)");
+  try {
+    const relPath = await resolveArticleRelPath(root, name); // throws if not found — validates `name`
+    const diff = await articleDiff(root, rev, relPath);
+    if (diff === null) return jsonErr("commit not found or no diff for this article at that revision", 404);
+    return jsonOk({ name, rev, diff });
+  } catch (e) {
+    return jsonErr((e as Error).message, 404);
+  }
+}
+
+// ---- POST /api/memory/init { path?: string } -----------------------------
+//
+// The consent-card backend: wraps the exact same `setupVault` the CLI's
+// `mlx-bun memory init` calls (src/cli.ts, `case "memory":` → `sub === "init"`)
+// — same idempotent create-dirs/README/Meta-pages/git-init behavior, minus
+// the CLI's interactive extras (seed-from-existing-vault prompt, nightly
+// schedule prompt), which stay CLI-only/TTY-gated and are out of scope here.
+//
+// Unlike the CLI (which only ever calls setupVault(vaultRoot()) — a
+// user-supplied path there is a read-only *seed source*, never the
+// setupVault target), this route takes `path` straight from an untrusted
+// loopback caller. setupVault() mkdir -p's subdirectories into it and, if
+// it isn't already a git repo, runs `git init && git add . && git commit`
+// there — pointed at an arbitrary directory that's a destructive foot-gun
+// (force-commits whatever's already there into a fresh history). There is
+// no product reason for the web panel to redirect the vault root outside
+// the machine's own tmp/vault trees, so `path` is confined to those before
+// it ever reaches setupVault; anything else 400s instead of writing files.
+
+/** realpath() of the deepest existing ancestor of `path`, with the
+ *  non-existent tail (e.g. a not-yet-created vault dir) re-appended
+ *  literally. Needed because macOS resolves `/var/folders/...` (os.tmpdir())
+ *  through a `/private` symlink: a brand-new mkdtemp path under it won't
+ *  literal-prefix-match the realpath'd tmpdir unless we walk up to a real
+ *  segment the same way. */
+async function realpathNearest(path: string): Promise<string> {
+  let cur = resolve(path);
+  const tail: string[] = [];
+  while (true) {
+    try {
+      return [await realpath(cur), ...tail.reverse()].join(sep);
+    } catch {
+      const parent = dirname(cur);
+      if (parent === cur) return resolve(path); // hit the fs root without finding anything real
+      tail.push(cur.slice(parent.length + 1));
+      cur = parent;
+    }
+  }
+}
+
+/** Is `candidate` equal to or a descendant of `ancestor`, after resolving
+ *  symlinks on both sides? */
+async function isUnder(candidate: string, ancestor: string): Promise<boolean> {
+  const realAncestor = await realpathNearest(ancestor);
+  const realCandidate = await realpathNearest(candidate);
+  return realCandidate === realAncestor || realCandidate.startsWith(realAncestor + sep);
+}
+
+/** Confine an init target to the machine's own vault/tmp trees — the only
+ *  legitimate destinations for this route (default vault root, or a tmp dir
+ *  for tests). Rejects anything else rather than handing setupVault() an
+ *  arbitrary, possibly pre-existing directory to mkdir/git-init into. */
+async function isAllowedVaultTarget(path: string): Promise<boolean> {
+  return (await isUnder(path, vaultRoot())) || (await isUnder(path, tmpdir()));
+}
+
+export async function handleMemoryInit(request: Request): Promise<Response> {
+  let body: { path?: string } = {};
+  try {
+    body = (await request.json()) as { path?: string };
+  } catch {
+    // empty body is fine — defaults to vaultRoot()
+  }
+  const root = body.path && body.path.trim() ? body.path.trim() : vaultRoot();
+  if (!(await isAllowedVaultTarget(root))) {
+    return jsonErr("path must be under the memory vault root or a temp directory", 400);
+  }
+  try {
+    const res = await setupVault(root);
+    const st = await vaultStatus(res.root);
+    return jsonOk({ result: res, status: st });
+  } catch (e) {
+    return jsonErr((e as Error).message, 500);
+  }
+}

package/src/memory/vault.ts

@@ -218,6 +218,17 @@ export async function readArticle(root: string, article: string): Promise<{ path
   return { path, content: await readFile(path, "utf8") };
 }
 
+/** Vault-relative path (`articles/<stem>.md` or `Reference/<stem>.md`) for an
+ *  existing article/reference doc — the safe argv piece for git plumbing
+ *  (articleHistory/articleDiff), resolved via the same name normalization/
+ *  existence check as readArticle so a caller can't smuggle a path-traversal
+ *  or shell-metacharacter name through. Throws if the doc doesn't exist. */
+export async function resolveArticleRelPath(root: string, article: string): Promise<string> {
+  const doc = normalizeMemoryDocId(article);
+  await resolveArticlePath(root, article); // existence check, reuses the same normalization
+  return doc.kind === "reference" ? join("Reference", `${doc.stem}.md`) : join("articles", `${doc.stem}.md`);
+}
+
 /** Filename stems for every `*.md` under articles/, sorted alphabetically. */
 export async function listArticles(root: string): Promise<string[]> {
   let names: string[];
@@ -1013,6 +1024,69 @@ function runGit(args: string[], cwd: string): Promise<boolean> {
   });
 }
 
+// ---- read-only git plumbing (per-article history/diff) ----------------
+//
+// Used by the memory REST surface (GET /api/memory/history|diff) to render
+// "watch it self-heal" — the article's commit history and a specific
+// commit's diff. Read-only plumbing only (log/show), args always an
+// explicit argv array via Bun.spawn (never a shell string), and callers
+// MUST validate `rev`/`name` before calling these (see server.ts) — this
+// layer trusts its `relPath`/`rev` inputs are already safe.
+
+export interface ArticleHistoryEntry {
+  hash: string;
+  date: string;
+  subject: string;
+}
+
+/** `git log` for one file, newest first, argv-only (never a shell string).
+ *  Returns [] when the vault isn't a git repo, the file has no history, or
+ *  git is unavailable. `relPath` must already be a repo-relative path (no
+ *  `..`, not absolute) — callers validate before calling. */
+export async function articleHistory(root: string, relPath: string, limit = 50): Promise<ArticleHistoryEntry[]> {
+  if (!(await pathExists(join(root, ".git")))) return [];
+  try {
+    const proc = Bun.spawn(
+      ["git", "log", `--max-count=${limit}`, "--date=short", "--pretty=format:%H%x1f%ad%x1f%s%x1e", "--", relPath],
+      { cwd: root, stdout: "pipe", stderr: "ignore" },
+    );
+    const out = await new Response(proc.stdout).text();
+    if ((await proc.exited) !== 0) return [];
+    return out
+      .split("\x1e")
+      .map((rec) => rec.trim())
+      .filter((rec) => rec.length > 0)
+      .map((rec) => {
+        const [hash = "", date = "", subject = ""] = rec.split("\x1f");
+        return { hash, date, subject };
+      })
+      .filter((e) => e.hash.length > 0);
+  } catch {
+    return [];
+  }
+}
+
+/** `git show <rev> -- <relPath>` — the diff for one commit touching one file.
+ *  `rev` must already be validated as a hex commit hash by the caller (never
+ *  a ref expression like `HEAD~1` or `main^`, which could be abused as an
+ *  argument-injection vector); `relPath` must already be repo-relative.
+ *  Returns null when the vault isn't a git repo or the commit/diff can't be
+ *  produced. */
+export async function articleDiff(root: string, rev: string, relPath: string): Promise<string | null> {
+  if (!(await pathExists(join(root, ".git")))) return null;
+  try {
+    const proc = Bun.spawn(
+      ["git", "show", rev, "--pretty=format:", "--", relPath],
+      { cwd: root, stdout: "pipe", stderr: "ignore" },
+    );
+    const out = await new Response(proc.stdout).text();
+    if ((await proc.exited) !== 0) return null;
+    return out;
+  } catch {
+    return null;
+  }
+}
+
 export interface SetupResult {
   root: string;
   created: string[];

package/src/mlx/array.ts

@@ -1,8 +1,8 @@
 // MlxArray: ownership wrapper around an mlx_array handle.
 // Explicit .dispose() is the contract; a FinalizationRegistry backstop
-// frees leaked handles on GC (verified in spikes/phase0-memory.ts).
+// frees leaked handles on GC (verified in lab/spikes/phase0-memory.ts).
 
-import { JSCallback, ptr, toArrayBuffer } from "bun:ffi";
+import { dlopen, ptr, toArrayBuffer } from "bun:ffi";
 import { C, Dtype, DTYPE_NAMES, type MlxHandle, optInt, outArray, takeMlxError } from "./ffi";
 import type { SafetensorsDtype } from "../safetensors";
 
@@ -24,17 +24,39 @@ export const SAFETENSORS_TO_MLX: Record<SafetensorsDtype, Dtype> = {
 
 // --- zero-copy pinning ---------------------------------------------------
 // Buffers handed to mlx via new_data_managed_payload must outlive the mlx
-// array. We pin the JS view in a map keyed by a payload id; the shared dtor
-// callback unpins when mlx drops its last reference. threadsafe: mlx may
-// release buffers from its eval thread.
+// array. fromView pins the JS view in this map (a GC root, nothing more);
+// release is EXPLICIT and JS-thread-only via unpinHostBuffer() — never an
+// mlx dtor. See the hazard note below.
 
 const pinned = new Map<number, Uint8Array>();
 let nextPinId = 1;
 
-const unpinCallback = new JSCallback(
-  (payloadId: number) => { pinned.delete(payloadId); },
-  { args: ["u64"], returns: "void", threadsafe: true },
-);
+// HAZARD (root cause of the 2026-07-06 restart-restore hang): mlx releases
+// buffer Data wherever the LAST shared_ptr drops — often the Metal
+// COMPLETION thread (gpu::eval retains buffers until the command buffer
+// finishes, i.e. past dispose()). A bun:ffi JSCallback dtor invoked from
+// that thread deadlocks when the JS thread is inside a blocking FFI eval
+// (completion waits on JS, JS waits on completion) and SIGTRAPs mid-GC.
+// So NO array here ever gets a JS dtor: fromPointer AND fromView both hand
+// mlx the native no-op dtor below (payload 0), and host-buffer lifetime is
+// the CALLER's contract, released only from the JS thread (unpinHostBuffer)
+// or never (process-lifetime mmaps: expert-offload).
+
+// Native no-op dtor: payload is always 0, and libc free(NULL) is defined to
+// do nothing — a dtor mlx can safely call from ANY thread. The pointed-to
+// memory's lifetime is the CALLER's contract (weight mmaps live for the
+// process; restored KV is COPIED at load since 2026-07-07 — see kv-store.ts
+// loadKvCache — so no restore mapping needs pinning anymore). The address
+// comes via dlsym because
+// bun:ffi's symbol .ptr returns the pointer bit-cast to float64
+// (lab/repro/bun-ffi-f64) — passing that back truncates to NULL.
+const libcDlsym = dlopen("/usr/lib/libSystem.B.dylib", {
+  dlsym: { args: ["u64", "ptr"], returns: "u64" },
+});
+const FREE_NAME = Buffer.from("free\0");
+const RTLD_DEFAULT = 0xfffffffffffffffen;
+const freeFnAddr = Number(libcDlsym.symbols.dlsym(RTLD_DEFAULT, ptr(FREE_NAME)));
+if (!freeFnAddr) throw new Error("dlsym(free) failed — zero-copy arrays need a native no-op dtor");
 
 export function pinnedBufferCount(): number {
   return pinned.size;
@@ -50,31 +72,61 @@ export class MlxArray {
   #handle: MlxHandle;
   #disposed = false;
   #token = {};
+  #pinId = 0; // fromView only: key into the pinned-view map (0 = no pin)
 
   constructor(handle: MlxHandle) {
     this.#handle = handle;
     registry.register(this, handle, this.#token);
   }
 
-  /** Zero-copy: wrap an existing buffer (e.g. an mmap view). The view is
-   *  pinned until mlx releases it. */
+  /** Zero-copy: wrap an existing JS buffer. The view is pinned (GC-rooted)
+   *  in a process-side map; mlx gets the NATIVE no-op dtor (free(NULL)) —
+   *  never a JS callback, so the last release may happen on the Metal
+   *  completion thread without deadlock (see hazard note above). Buffer
+   *  lifetime is therefore the CALLER's contract:
+   *   - mmap-backed views (toArrayBuffer aliases): the mapping owns the
+   *     memory; prefer fromPointer, which skips the pin entirely.
+   *   - JS-heap buffers: the pin keeps the buffer alive; call
+   *     unpinHostBuffer() only after a JS-side sync point proves mlx is
+   *     done with it (everything depending on the array evaluated AND the
+   *     array disposed), or never (process-lifetime, e.g. short scripts). */
   static fromView(view: Uint8Array, shape: number[], dtype: Dtype): MlxArray {
     const sb = shapeBuf(shape);
     const id = nextPinId++;
     pinned.set(id, view);
     const handle = C.mlx_array_new_data_managed_payload(
-      ptr(view), ptr(sb), shape.length, dtype, id, unpinCallback.ptr,
+      ptr(view), ptr(sb), shape.length, dtype, 0, freeFnAddr,
     );
-    return new MlxArray(handle);
+    const arr = new MlxArray(handle);
+    arr.#pinId = id;
+    return arr;
+  }
+
+  /** Release the GC root fromView placed on this array's host buffer.
+   *  JS-thread-only by construction (it's a plain method) — this is the
+   *  replacement for the deadlocking JSCallback unpin dtor. Idempotent;
+   *  no-op for non-fromView arrays. Caller contract: only call once mlx can
+   *  no longer touch the buffer (see fromView doc), or when the memory is
+   *  owned elsewhere (an mmap kept mapped past every use). */
+  unpinHostBuffer(): void {
+    if (this.#pinId !== 0) {
+      pinned.delete(this.#pinId);
+      this.#pinId = 0;
+    }
   }
 
   /** Zero-copy: wrap a raw pointer (e.g. into an mmap'd weight shard).
-   *  The caller guarantees the memory outlives the array — weight mmaps
-   *  live for the process, so the dtor is a no-op unpin of id 0. */
+   *  The caller guarantees the memory outlives EVERY mlx
+   *  reference to the array — including GPU command buffers that retain the
+   *  buffer past dispose() — i.e. mmaps backing these arrays stay mapped for
+   *  the process. The dtor is free(NULL) (payload 0): native + no-op, so
+   *  mlx may run it from the Metal completion thread without touching the
+   *  JS runtime (the old JSCallback dtor deadlocked serving / SIGTRAPed
+   *  under GC when that thread held the last reference — 2026-07-06). */
   static fromPointer(dataPtr: number, shape: number[], dtype: Dtype): MlxArray {
     const sb = shapeBuf(shape);
     const handle = C.mlx_array_new_data_managed_payload(
-      dataPtr, ptr(sb), shape.length, dtype, 0, unpinCallback.ptr,
+      dataPtr, ptr(sb), shape.length, dtype, 0, freeFnAddr,
     );
     return new MlxArray(handle);
   }
@@ -138,7 +190,7 @@ export class MlxArray {
     return Number(C.mlx_array_ndim(this.handle));
   }
 
-  // NOTE on the DFG stale-read bug (repro/bun-ffi-f64/ISSUE.md): the
+  // NOTE on the DFG stale-read bug (lab/repro/bun-ffi-f64/ISSUE.md): the
   // toArrayBuffer readbacks below (shape/rawBytes/toFloat32) are safe.
   // The hazard is reading a *pre-existing* typed array after an FFI call
   // wrote through its pointer — the JIT forwards a stale value from before
@@ -183,6 +235,18 @@ export class MlxArray {
 
   /** Raw bytes of the evaluated array (copy). bf16/f16/f32 only. */
   rawBytes(): Uint8Array {
+    return this.rawBytesView().slice();
+  }
+
+  /** ZERO-COPY view of the evaluated array's bytes — aliases the mlx
+   *  buffer directly, no JS-heap copy. Valid only while THIS array is
+   *  alive and unmutated: do not retain past dispose(). Built for
+   *  bulk-write paths (kv-store's streamed writer) where per-tensor
+   *  rawBytes() copies pile up in the JS heap faster than GC collects
+   *  them — the residual A7 whole-entry RSS spike (2026-07-07). The
+   *  array must be contiguous (ops.contiguous first for views), same
+   *  contract as rawBytes. */
+  rawBytesView(): Uint8Array {
     this.eval();
     const dt = this.dtype;
     const p =
@@ -190,9 +254,32 @@ export class MlxArray {
       : dt === Dtype.float16 ? C.mlx_array_data_float16(this.handle)
       : dt === Dtype.bfloat16 ? C.mlx_array_data_bfloat16(this.handle)
       : dt === Dtype.uint32 ? C.mlx_array_data_uint32(this.handle)
+      : dt === Dtype.uint8 ? C.mlx_array_data_uint8(this.handle)
+      : dt === Dtype.int8 ? C.mlx_array_data_int8(this.handle)
       : null;
-    if (p === null) throw new Error(`rawBytes: unsupported dtype ${this.dtypeName}`);
-    return new Uint8Array(toArrayBuffer(p!, 0, this.nbytes)).slice();
+    if (p === null) throw new Error(`rawBytesView: unsupported dtype ${this.dtypeName}`);
+    return new Uint8Array(toArrayBuffer(p!, 0, this.nbytes));
+  }
+
+  /** Read back an INTEGER array without enqueueing a cast kernel: eval +
+   *  direct buffer read for uint32/int32. A cast (`astype`) is a NEW op that
+   *  queues BEHIND everything already dispatched on the GPU stream — on the
+   *  batched decode hot path, reading the pipeline register via toFloat32
+   *  stalled the "overlapped" token read for a FULL step of GPU work every
+   *  token (the core of the Phase-0 B=1 gap; unified-engine plan Phase 2).
+   *  Non-integer dtypes fall back to the rounded float readback. */
+  toIntTokens(): number[] {
+    if (this.dtype === Dtype.uint32) {
+      this.eval();
+      const p = C.mlx_array_data_uint32(this.handle);
+      return [...new Uint32Array(toArrayBuffer(p!, 0, this.size * 4))];
+    }
+    if (this.dtype === Dtype.int32) {
+      this.eval();
+      const p = C.mlx_array_data_int32(this.handle);
+      return [...new Int32Array(toArrayBuffer(p!, 0, this.size * 4))];
+    }
+    return [...this.toFloat32()].map((x) => Math.round(x));
   }
 
   /** Read back as float32 (casts on GPU if needed, then copies out). */
@@ -205,6 +292,49 @@ export class MlxArray {
     return out;
   }
 
+  /** Read an f32/f16/bf16 array numerically as float32 without enqueueing a
+   *  GPU cast. This evaluates the ORIGINAL array, then expands 16-bit values
+   *  on the host. Use this for small readbacks after later GPU work has
+   *  already been dispatched: toFloat32() would append astype(float32) behind
+   *  that work and turn the readback into an accidental full-step barrier. */
+  toFloat32Host(): Float32Array {
+    this.eval();
+    const dt = this.dtype;
+    if (dt === Dtype.float32) {
+      const p = C.mlx_array_data_float32(this.handle);
+      return new Float32Array(toArrayBuffer(p!, 0, this.size * 4)).slice();
+    }
+    if (dt === Dtype.bfloat16) {
+      const p = C.mlx_array_data_bfloat16(this.handle);
+      const src = new Uint16Array(toArrayBuffer(p!, 0, this.size * 2));
+      const words = new Uint32Array(src.length);
+      for (let i = 0; i < src.length; i++) words[i] = src[i]! << 16;
+      return new Float32Array(words.buffer);
+    }
+    if (dt === Dtype.float16) {
+      const p = C.mlx_array_data_float16(this.handle);
+      const src = new Uint16Array(toArrayBuffer(p!, 0, this.size * 2));
+      const out = new Float32Array(src.length);
+      for (let i = 0; i < src.length; i++) {
+        const bits = src[i]!;
+        const sign = bits & 0x8000;
+        const exponent = (bits >>> 10) & 0x1f;
+        const fraction = bits & 0x03ff;
+        let value: number;
+        if (exponent === 0) {
+          value = fraction === 0 ? 0 : fraction * 2 ** -24;
+        } else if (exponent === 0x1f) {
+          value = fraction === 0 ? Infinity : NaN;
+        } else {
+          value = (1 + fraction / 1024) * 2 ** (exponent - 15);
+        }
+        out[i] = sign ? -value : value;
+      }
+      return out;
+    }
+    throw new Error(`toFloat32Host: unsupported dtype ${this.dtypeName}`);
+  }
+
   /** Basic slice: start/stop per dimension (stride 1). */
   slice(start: number[], stop: number[], stream: MlxHandle = gpuStream): MlxArray {
     const s = new Int32Array(start);