@mlx-node/agent 0.0.8 → 0.0.10

package/dist/provider/stream-adapter.js

@@ -51,6 +51,16 @@ function safeString(read, fallback) {
         return fallback;
     }
 }
+/** Read numeric model metadata without trusting a user-configurable getter. */
+function readPositiveSafeInteger(read) {
+    try {
+        const value = read();
+        return typeof value === 'number' && Number.isSafeInteger(value) && value > 0 ? value : undefined;
+    }
+    catch {
+        return undefined;
+    }
+}
 /**
  * Publish the native model's load-time physical context limit onto pi's shared
  * model object. The parent session and every in-process subagent resolve this
@@ -62,13 +72,18 @@ function safeString(read, fallback) {
  * discovery-time limit here, and never let metadata synchronization break an
  * otherwise valid inference turn.
  */
-function publishEffectiveContextWindow(model, session) {
+function publishEffectiveContextWindow(model, session, configuredModelMaxTokens) {
     try {
         const effective = Math.floor(session.contextLimits()?.effectiveWindowTokens ?? 0);
         if (!Number.isSafeInteger(effective) || effective <= 0)
             return;
         model.contextWindow = Math.min(model.contextWindow, effective);
-        model.maxTokens = Math.min(model.maxTokens, model.contextWindow);
+        // Use the already-validated snapshot. Reading the shared metadata here
+        // would let Math.min coerce invalid user configuration into a valid value
+        // that a later turn would then trust.
+        if (configuredModelMaxTokens !== undefined) {
+            model.maxTokens = Math.min(configuredModelMaxTokens, model.contextWindow);
+        }
     }
     catch {
         // Exact native preflight still protects capacity; keep serving the turn.
@@ -125,7 +140,7 @@ function failsafeMessage(model, reason, message) {
         timestamp: Date.now(),
     };
 }
-export function makeMlxStreamSimple(host, onPerformance, resolveRootCacheOwner) {
+export function makeMlxStreamSimple(host, onPerformance, resolveRootCacheOwner, onTurnRecord, onTurnStart, resolveRootSessionFile) {
     return (model, context, options) => {
         const stream = createAssistantMessageEventStream();
         /**
@@ -138,6 +153,7 @@ export function makeMlxStreamSimple(host, onPerformance, resolveRootCacheOwner)
         let emitter;
         let signal;
         let rootCacheOwnerId;
+        let rootSessionFile;
         let resolvedReasoning;
         let detachAbort;
         /**
@@ -208,8 +224,11 @@ export function makeMlxStreamSimple(host, onPerformance, resolveRootCacheOwner)
             signal = options?.signal;
             // Snapshot the top-level owner before this request can queue behind
             // another inference. A later /new or /resume must not relabel an older
-            // request that was already submitted under the previous root.
+            // request that was already submitted under the previous root. The metrics
+            // root id reuses this same session_start value; only the JSONL path is a
+            // separate snapshot, taken here so it can never drift from the id.
             rootCacheOwnerId = resolveRootCacheOwner?.();
+            rootSessionFile = resolveRootSessionFile?.();
             // Snapshot once: the native config and the replay provenance must describe
             // the same resolved template mode. Presence alone is wrong for Pi's
             // minimal/low levels, both of which resolve to disabled thinking.
@@ -238,13 +257,42 @@ export function makeMlxStreamSimple(host, onPerformance, resolveRootCacheOwner)
         }
         void (async () => {
             let sawNativeFinal = false;
-            await host.runWithResident(model.id, async (session) => {
+            // Bracket the whole serialized turn (queue wait + resident selection +
+            // prefill + decode) so a MetricsTrace record can report wall-clock
+            // duration. Read at the terminal only on the success path below.
+            const turnStartedAt = Date.now();
+            await host.runWithResident(model.id, async (session, resident) => {
+                // Callback entry: the queue wait + any cold model load/swap have now
+                // resolved, but no native work (prime/prefill/decode) has started. This
+                // is the seam that separates queue+load latency from execution time —
+                // `durationMs` (read at the terminal) still brackets the whole turn.
+                const execStartedAt = Date.now();
+                const queueMs = Math.max(0, execStartedAt - turnStartedAt);
                 // Terminated while queued behind earlier inference/loading (or
                 // between stages below): the terminal already went out — skip ALL
                 // session work (no warm-reset, no prime, no stream).
                 if (terminated)
                     return;
-                publishEffectiveContextWindow(model, session);
+                // Serialized-turn start: fire before any native work so a metrics
+                // consumer can snapshot process-wide cold-tier counters here and diff
+                // them at the success terminal. Because turns are serialized, a prior
+                // turn that aborted/errored has already fully drained by now, so its
+                // activity is baked into this snapshot and never attributed to this
+                // turn's delta. Best-effort — a throwing hook must not break inference.
+                if (onTurnStart) {
+                    try {
+                        onTurnStart();
+                    }
+                    catch {
+                        // Telemetry snapshot is best-effort; never break inference.
+                    }
+                }
+                // Snapshot the composed budget before publishing native context limits.
+                // The advisory publisher must not read this raw value itself:
+                // Math.min can coerce hostile metadata (for example "512" or Infinity)
+                // into a finite number that a later turn would then trust.
+                const configuredModelMaxTokens = readPositiveSafeInteger(() => model.maxTokens);
+                publishEffectiveContextWindow(model, session, configuredModelMaxTokens);
                 const supportsImages = publishImageCapability(model, session);
                 const discovered = host.modelInfo(model.id);
                 if (!discovered) {
@@ -271,7 +319,7 @@ export function makeMlxStreamSimple(host, onPerformance, resolveRootCacheOwner)
                     return;
                 try {
                     session.primeHistory(contextToChatMessages(context, supportsImages));
-                    const config = buildChatConfig(discovered.modelType, options, toolsToDefinitions(context.tools), rootCacheOwnerId, resolvedReasoning);
+                    const config = buildChatConfig(discovered.modelType, options, toolsToDefinitions(context.tools), rootCacheOwnerId, resolvedReasoning, configuredModelMaxTokens);
                     for await (const event of session.startFromHistoryStream(config, signal)) {
                         if (event.done) {
                             if (event.finishReason === 'error') {
@@ -308,6 +356,33 @@ export function makeMlxStreamSimple(host, onPerformance, resolveRootCacheOwner)
                                     detachAbort = undefined;
                                     try {
                                         turn.onFinal(event);
+                                        // Durable per-turn telemetry: only a successful final gets
+                                        // here. Guarded independently of onFinal so a throwing
+                                        // recorder cannot derail the (already-emitted) terminal.
+                                        if (onTurnRecord) {
+                                            try {
+                                                onTurnRecord({
+                                                    traceId: turn.traceId,
+                                                    sessionId: options?.sessionId,
+                                                    // Root correlation is the SUBMIT-time snapshot, not the
+                                                    // live root, so a /new or /resume that landed while this
+                                                    // turn ran can't reattribute it (id == cacheOwner root).
+                                                    rootSessionId: rootCacheOwnerId,
+                                                    rootSessionFile,
+                                                    model: model.id,
+                                                    final: event,
+                                                    durationMs: Math.max(0, Date.now() - turnStartedAt),
+                                                    // Queue+load wait captured at callback entry, and the
+                                                    // warm/cold distinction from `runWithResident` — so the
+                                                    // record can attribute wait vs execution latency.
+                                                    queueMs,
+                                                    resident,
+                                                });
+                                            }
+                                            catch {
+                                                // Telemetry is best-effort; never break inference.
+                                            }
+                                        }
                                     }
                                     catch (err) {
                                         pushFailsafeTerminal('error', coerceErrorMessage(err));

package/dist/provider/warm-reuse.d.ts

@@ -18,18 +18,19 @@
  * the reused prefix and skip the corresponding re-prefill.
  *
  * Fields accessed: `inFlight`, `history`, `lastImagesKey`, `lastAudioKey`, `turnCount`,
- * `unresolvedOkToolCallCount`, `needsFullReplay`. These are TypeScript `private` fields on
- * `ChatSession` (compile-time only) — at runtime they are ordinary
- * properties. The cast through {@link ChatSessionWarmReuseInternals}
- * gives this helper a typed view of the instance without relaxing the
- * class's `private` declarations. The field names MUST stay in sync
- * with `packages/lm/src/chat-session.ts`; a mismatch would silently
- * skip the intended state wipe — the drift test in
+ * `unresolvedOkToolCallCount`, `needsFullReplay`, `defaultConfig`, `activeTools`.
+ * These are TypeScript `private` fields on `ChatSession` (compile-time
+ * only) — at runtime they are ordinary properties. The cast through
+ * {@link ChatSessionWarmReuseInternals} gives this helper a typed view
+ * of the instance without relaxing the class's `private` declarations.
+ * The field names MUST stay in sync with
+ * `packages/lm/src/chat-session.ts`; a mismatch would silently skip
+ * the intended state wipe — the drift test in
  * `packages/agent/__test__/warm-reuse.test.ts` checks every name in
  * {@link WARM_REUSE_TOUCHED_FIELDS} against a real `ChatSession`
  * instance.
  */
-import type { ChatSession, SessionCapableModel } from '@mlx-node/lm';
+import type { ChatConfig, ChatSession, SessionCapableModel } from '@mlx-node/lm';
 /**
  * Private structural view of the `ChatSession` JS-side state that the
  * warm-reuse helper needs to wipe. Mirrors the internal state
@@ -45,6 +46,8 @@ interface ChatSessionWarmReuseInternals {
     turnCount: number;
     unresolvedOkToolCallCount: number | null;
     needsFullReplay: boolean;
+    defaultConfig?: ChatConfig;
+    activeTools: ChatConfig['tools'];
 }
 /**
  * Runtime list of the `ChatSession` private field names this module

package/dist/provider/warm-reuse.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"warm-reuse.d.ts","sourceRoot":"","sources":["../../src/provider/warm-reuse.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;GA8BG;AAEH,OAAO,KAAK,EAAE,WAAW,EAAE,mBAAmB,EAAE,MAAM,cAAc,CAAC;AAErE;;;;;;GAMG;AACH,UAAU,6BAA6B;IACrC,QAAQ,EAAE,OAAO,CAAC;IAClB,OAAO,EAAE,OAAO,EAAE,CAAC;IACnB,aAAa,EAAE,MAAM,GAAG,IAAI,CAAC;IAC7B,YAAY,EAAE,MAAM,GAAG,IAAI,CAAC;IAC5B,SAAS,EAAE,MAAM,CAAC;IAClB,yBAAyB,EAAE,MAAM,GAAG,IAAI,CAAC;IACzC,eAAe,EAAE,OAAO,CAAC;CAC1B;AAkBD;;;;GAIG;AACH,eAAO,MAAM,yBAAyB,EAAgD,aAAa,CACjG,MAAM,6BAA6B,CACpC,CAAC;AAEF;;;;;;;;;;;;;;;GAeG;AACH,wBAAsB,sCAAsC,CAAC,CAAC,SAAS,mBAAmB,EACxF,OAAO,EAAE,WAAW,CAAC,CAAC,CAAC,GACtB,OAAO,CAAC,IAAI,CAAC,CAqBf"}
+{"version":3,"file":"warm-reuse.d.ts","sourceRoot":"","sources":["../../src/provider/warm-reuse.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;GA+BG;AAEH,OAAO,KAAK,EAAE,UAAU,EAAE,WAAW,EAAE,mBAAmB,EAAE,MAAM,cAAc,CAAC;AAEjF;;;;;;GAMG;AACH,UAAU,6BAA6B;IACrC,QAAQ,EAAE,OAAO,CAAC;IAClB,OAAO,EAAE,OAAO,EAAE,CAAC;IACnB,aAAa,EAAE,MAAM,GAAG,IAAI,CAAC;IAC7B,YAAY,EAAE,MAAM,GAAG,IAAI,CAAC;IAC5B,SAAS,EAAE,MAAM,CAAC;IAClB,yBAAyB,EAAE,MAAM,GAAG,IAAI,CAAC;IACzC,eAAe,EAAE,OAAO,CAAC;IACzB,aAAa,CAAC,EAAE,UAAU,CAAC;IAC3B,WAAW,EAAE,UAAU,CAAC,OAAO,CAAC,CAAC;CAClC;AAoBD;;;;GAIG;AACH,eAAO,MAAM,yBAAyB,EAAgD,aAAa,CACjG,MAAM,6BAA6B,CACpC,CAAC;AAEF;;;;;;;;;;;;;;;GAeG;AACH,wBAAsB,sCAAsC,CAAC,CAAC,SAAS,mBAAmB,EACxF,OAAO,EAAE,WAAW,CAAC,CAAC,CAAC,GACtB,OAAO,CAAC,IAAI,CAAC,CAyBf"}

package/dist/provider/warm-reuse.js

@@ -18,13 +18,14 @@
  * the reused prefix and skip the corresponding re-prefill.
  *
  * Fields accessed: `inFlight`, `history`, `lastImagesKey`, `lastAudioKey`, `turnCount`,
- * `unresolvedOkToolCallCount`, `needsFullReplay`. These are TypeScript `private` fields on
- * `ChatSession` (compile-time only) — at runtime they are ordinary
- * properties. The cast through {@link ChatSessionWarmReuseInternals}
- * gives this helper a typed view of the instance without relaxing the
- * class's `private` declarations. The field names MUST stay in sync
- * with `packages/lm/src/chat-session.ts`; a mismatch would silently
- * skip the intended state wipe — the drift test in
+ * `unresolvedOkToolCallCount`, `needsFullReplay`, `defaultConfig`, `activeTools`.
+ * These are TypeScript `private` fields on `ChatSession` (compile-time
+ * only) — at runtime they are ordinary properties. The cast through
+ * {@link ChatSessionWarmReuseInternals} gives this helper a typed view
+ * of the instance without relaxing the class's `private` declarations.
+ * The field names MUST stay in sync with
+ * `packages/lm/src/chat-session.ts`; a mismatch would silently skip
+ * the intended state wipe — the drift test in
  * `packages/agent/__test__/warm-reuse.test.ts` checks every name in
  * {@link WARM_REUSE_TOUCHED_FIELDS} against a real `ChatSession`
  * instance.
@@ -43,6 +44,8 @@ const WARM_REUSE_TOUCHED_FIELD_SET = {
     turnCount: true,
     unresolvedOkToolCallCount: true,
     needsFullReplay: true,
+    defaultConfig: true,
+    activeTools: true,
 };
 /**
  * Runtime list of the `ChatSession` private field names this module
@@ -85,4 +88,8 @@ export async function resetPreservingNativeCacheForWarmReuse(session) {
     internals.turnCount = 0;
     internals.unresolvedOkToolCallCount = null;
     internals.needsFullReplay = false;
+    // Tools are conversation state. A provider replay can switch to an
+    // unrelated history, so restore constructor defaults exactly like
+    // ChatSession.reset() instead of leaking the prior committed overlay.
+    internals.activeTools = internals.defaultConfig?.tools;
 }

package/dist/run-agent.d.ts

@@ -17,7 +17,7 @@
  *   critical on pi's hard `process.exit()` paths.
  */
 import type { InlineExtension } from '@earendil-works/pi-coding-agent';
-import { type FilterableModelRegistryConstructor } from './provider/model-registry-filter.js';
+import { type FilterableModelRuntimeConstructor } from './provider/model-registry-filter.js';
 import type { MlxModelInfo } from './provider/models.js';
 /** Shape of pi's `main(argv, { extensionFactories })` — also the test seam. */
 export type RunAgentMain = (args: string[], opts: {
@@ -25,11 +25,11 @@ export type RunAgentMain = (args: string[], opts: {
 }) => Promise<void>;
 export interface RunAgentPi {
     main: RunAgentMain;
-    ModelRegistry: FilterableModelRegistryConstructor;
+    ModelRuntime: FilterableModelRuntimeConstructor;
 }
 /** @internal Narrow lifecycle seam for the agent's paged config overlays. */
 export interface AgentPagedConfigOverrides {
-    resolve(modelPath: string, modelType?: string): Promise<string>;
+    resolve(modelPath: string, modelType?: string, persistPagedCache?: boolean): Promise<string>;
     cleanup(): Promise<void>;
 }
 export interface RunAgentOptions {
@@ -41,6 +41,15 @@ export interface RunAgentOptions {
     argv: string[];
     /** Native inference-log path to surface after Pi takes over the TUI. */
     traceLogFile?: string;
+    /**
+     * Enable the SSD cold tier by default (the agent's default; the CLI sets it
+     * false for `--no-persist-cache`). Forwarded to {@link MlxModelHost}, which
+     * applies this ONE value to every load whose family is in
+     * `COLD_TIER_RESTORE_FAMILIES` — not to qwen3 alone. Families off that list
+     * are handed no policy because they can never persist, not because this flag
+     * spares them. `undefined` keeps the host's on-by-default behavior.
+     */
+    persistPagedCache?: boolean;
     /** @internal Test seam; when set, the pi dynamic import is skipped entirely. */
     piImpl?: RunAgentPi;
     /** @internal Test seam for paged model-path resolution and cleanup. */

package/dist/run-agent.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"run-agent.d.ts","sourceRoot":"","sources":["../src/run-agent.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;;;;;;;;GAiBG;AAKH,OAAO,KAAK,EAAE,eAAe,EAAE,MAAM,iCAAiC,CAAC;AASvE,OAAO,EACL,KAAK,kCAAkC,EAExC,MAAM,qCAAqC,CAAC;AAC7C,OAAO,KAAK,EAAE,YAAY,EAAE,MAAM,sBAAsB,CAAC;AAEzD,+EAA+E;AAC/E,MAAM,MAAM,YAAY,GAAG,CAAC,IAAI,EAAE,MAAM,EAAE,EAAE,IAAI,EAAE;IAAE,kBAAkB,EAAE,eAAe,EAAE,CAAA;CAAE,KAAK,OAAO,CAAC,IAAI,CAAC,CAAC;AAE9G,MAAM,WAAW,UAAU;IACzB,IAAI,EAAE,YAAY,CAAC;IACnB,aAAa,EAAE,kCAAkC,CAAC;CACnD;AAED,6EAA6E;AAC7E,MAAM,WAAW,yBAAyB;IACxC,OAAO,CAAC,SAAS,EAAE,MAAM,EAAE,SAAS,CAAC,EAAE,MAAM,GAAG,OAAO,CAAC,MAAM,CAAC,CAAC;IAChE,OAAO,IAAI,OAAO,CAAC,IAAI,CAAC,CAAC;CAC1B;AAED,MAAM,WAAW,eAAe;IAC9B,2FAA2F;IAC3F,SAAS,EAAE,MAAM,CAAC;IAClB,wEAAwE;IACxE,MAAM,EAAE,YAAY,EAAE,CAAC;IACvB,yDAAyD;IACzD,IAAI,EAAE,MAAM,EAAE,CAAC;IACf,wEAAwE;IACxE,YAAY,CAAC,EAAE,MAAM,CAAC;IACtB,gFAAgF;IAChF,MAAM,CAAC,EAAE,UAAU,CAAC;IACpB,uEAAuE;IACvE,oBAAoB,CAAC,EAAE,yBAAyB,CAAC;CAClD;AAED,mFAAmF;AACnF,wBAAgB,sBAAsB,CAAC,GAAG,GAAE,MAAM,CAAC,UAAwB,GAAG,OAAO,CAEpF;AAED;;;;;;;;;GASG;AACH,wBAAsB,QAAQ,CAAC,IAAI,EAAE,eAAe,GAAG,OAAO,CAAC,IAAI,CAAC,CAmDnE"}
+{"version":3,"file":"run-agent.d.ts","sourceRoot":"","sources":["../src/run-agent.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;;;;;;;;GAiBG;AAKH,OAAO,KAAK,EAAE,eAAe,EAAE,MAAM,iCAAiC,CAAC;AAWvE,OAAO,EACL,KAAK,iCAAiC,EAEvC,MAAM,qCAAqC,CAAC;AAC7C,OAAO,KAAK,EAAE,YAAY,EAAE,MAAM,sBAAsB,CAAC;AAEzD,+EAA+E;AAC/E,MAAM,MAAM,YAAY,GAAG,CAAC,IAAI,EAAE,MAAM,EAAE,EAAE,IAAI,EAAE;IAAE,kBAAkB,EAAE,eAAe,EAAE,CAAA;CAAE,KAAK,OAAO,CAAC,IAAI,CAAC,CAAC;AAE9G,MAAM,WAAW,UAAU;IACzB,IAAI,EAAE,YAAY,CAAC;IACnB,YAAY,EAAE,iCAAiC,CAAC;CACjD;AAED,6EAA6E;AAC7E,MAAM,WAAW,yBAAyB;IACxC,OAAO,CAAC,SAAS,EAAE,MAAM,EAAE,SAAS,CAAC,EAAE,MAAM,EAAE,iBAAiB,CAAC,EAAE,OAAO,GAAG,OAAO,CAAC,MAAM,CAAC,CAAC;IAC7F,OAAO,IAAI,OAAO,CAAC,IAAI,CAAC,CAAC;CAC1B;AAED,MAAM,WAAW,eAAe;IAC9B,2FAA2F;IAC3F,SAAS,EAAE,MAAM,CAAC;IAClB,wEAAwE;IACxE,MAAM,EAAE,YAAY,EAAE,CAAC;IACvB,yDAAyD;IACzD,IAAI,EAAE,MAAM,EAAE,CAAC;IACf,wEAAwE;IACxE,YAAY,CAAC,EAAE,MAAM,CAAC;IACtB;;;;;;;OAOG;IACH,iBAAiB,CAAC,EAAE,OAAO,CAAC;IAC5B,gFAAgF;IAChF,MAAM,CAAC,EAAE,UAAU,CAAC;IACpB,uEAAuE;IACvE,oBAAoB,CAAC,EAAE,yBAAyB,CAAC;CAClD;AAED,mFAAmF;AACnF,wBAAgB,sBAAsB,CAAC,GAAG,GAAE,MAAM,CAAC,UAAwB,GAAG,OAAO,CAEpF;AAED;;;;;;;;;GASG;AACH,wBAAsB,QAAQ,CAAC,IAAI,EAAE,eAAe,GAAG,OAAO,CAAC,IAAI,CAAC,CA2EnE"}

package/dist/run-agent.js

@@ -18,7 +18,9 @@
  */
 import { homedir } from 'node:os';
 import { join } from 'node:path';
+import { coldCacheDrain } from '@mlx-node/core';
 import { PagedConfigOverrideManager } from '@mlx-node/lm';
+import { createLocalImageInputExtension } from './extensions/local-image-input.js';
 import { createPermissionGateExtension } from './extensions/permission-gate.js';
 import { createSubagentExtension } from './extensions/subagent.js';
 import { createTerminalTitleExtension } from './extensions/terminal-title.js';
@@ -46,6 +48,15 @@ export async function runAgent(opts) {
     // Mirrors `mlx launch claude`: chunked paged prefill keeps long-prompt
     // TTFT bounded on the default paged path.
     process.env.MLX_PAGED_PREFILL_CHUNK_SIZE ??= '2048';
+    // Hard offline invariant — NOT a user-overridable default, hence `=` not `??=`.
+    // `mlx agent` is local-only: no cloud provider may ever be contacted. pi 0.81.1's
+    // interactive and RPC startup call `ModelRuntime.refresh()`, which when PI_OFFLINE
+    // is unset fetches remote provider catalogs from pi.dev and can refresh a persisted
+    // cloud credential — a network path the mlx-only prototype filter does NOT cover
+    // (it patches the read methods, not `refresh`). Forcing PI_OFFLINE=1 here, before pi
+    // is imported below, also pins every ModelRuntime in this process to `allowNetwork`
+    // off, so no ambient/prior cloud credential can leak outbound traffic.
+    process.env.PI_OFFLINE = '1';
     // Force every paged-capable agent family through an isolated config clone.
     // This includes quantized LFM2 (whose standalone default is deliberately
     // flat) and Qwen3.5 dense/MoE (whose text-only defaults are flat). Gemma4's
@@ -55,19 +66,23 @@ export async function runAgent(opts) {
     const preserveEmbeddedGemmaDraft = agentGemmaDraftEnabled();
     const pagedConfigOverrides = opts.pagedConfigOverrides ?? new PagedConfigOverrideManager({ preserveEmbeddedGemmaDraft });
     const modelHost = new MlxModelHost(opts.models.map((model) => model.discovered), {
-        resolveModelPathFn: (model) => pagedConfigOverrides.resolve(model.path, model.modelType),
+        resolveModelPathFn: (model, policy) => pagedConfigOverrides.resolve(model.path, model.modelType, policy?.persistPagedCache),
         requirePagedCache: true,
+        persistPagedCache: opts.persistPagedCache,
     });
-    // Keep the pi import strictly behind the seam. The seam carries BOTH main
-    // and its registry class, so tests and production exercise the same policy
-    // installation/lifecycle instead of being able to bypass it accidentally.
+    // Keep the pi import strictly behind the seam. The seam carries BOTH main and
+    // the ModelRuntime class, so tests and production exercise the same policy
+    // installation/lifecycle instead of being able to bypass it accidentally. The
+    // filter patches the runtime prototype (not the extension-only ModelRegistry
+    // facade), which is where the selector / listing / resolution paths read.
     const pi = opts.piImpl ?? (await import('@earendil-works/pi-coding-agent'));
-    const restoreModelRegistry = installMlxOnlyModelRegistryFilter(pi.ModelRegistry, opts.models.map((model) => model.discovered.name));
+    const restoreModelRegistry = installMlxOnlyModelRegistryFilter(pi.ModelRuntime, opts.models.map((model) => model.discovered.name));
     const subagentsEnabled = opts.models.length > 0 && !opts.argv.includes('--no-extensions') && !opts.argv.includes('-ne');
     try {
         await pi.main(opts.argv, {
             extensionFactories: [
                 createMlxProviderExtension(opts.models, modelHost),
+                createLocalImageInputExtension(),
                 createPermissionGateExtension(),
                 ...(subagentsEnabled ? [createSubagentExtension()] : []),
                 ...(opts.traceLogFile !== undefined ? [createTraceNoticeExtension(opts.traceLogFile)] : []),
@@ -81,6 +96,17 @@ export async function runAgent(opts) {
         }
         finally {
             await pagedConfigOverrides.cleanup();
+            // `mlx agent -p` is one-shot: flush any accepted cold-tier prefix blocks
+            // to disk before the process exits, otherwise a prompt's just-persisted
+            // KV could still be queued/mid-write when we return. No-op when the tier
+            // was never opened, bounded so a stuck fsync can't hang exit, and never
+            // allowed to throw out of cleanup (best-effort durability).
+            try {
+                coldCacheDrain(5000);
+            }
+            catch {
+                // Best-effort: a drain failure must never mask the real exit path.
+            }
         }
     }
 }

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@mlx-node/agent",
-  "version": "0.0.8",
+  "version": "0.0.10",
   "homepage": "https://github.com/mlx-node/mlx-node",
   "bugs": {
     "url": "https://github.com/mlx-node/mlx-node/issues"
@@ -21,16 +21,21 @@
     ".": {
       "types": "./dist/index.d.ts",
       "import": "./dist/index.js"
+    },
+    "./catalog": {
+      "types": "./dist/catalog.d.ts",
+      "default": "./dist/catalog.js"
     }
   },
   "scripts": {
     "build": "tsc -b"
   },
   "dependencies": {
-    "@earendil-works/pi-ai": "0.80.6",
-    "@earendil-works/pi-coding-agent": "0.80.6",
-    "@mlx-node/lm": "0.0.8",
-    "@mlx-node/server": "0.0.8",
+    "@earendil-works/pi-ai": "0.81.1",
+    "@earendil-works/pi-coding-agent": "0.81.1",
+    "@mlx-node/core": "0.0.10",
+    "@mlx-node/lm": "0.0.10",
+    "@mlx-node/server": "0.0.10",
     "typebox": "1.3.6"
   },
   "devDependencies": {