@mlx-node/agent 0.0.8 → 0.0.9

package/dist/provider/model-registry-filter.js

@@ -1,72 +1,263 @@
 /**
- * Process-local policy adapter for pi's public ModelRegistry.
+ * Process-local policy adapter for pi's canonical `ModelRuntime`.
  *
- * `mlx agent` is an offline/local product, but pi's registry also contains
- * every authenticated built-in cloud provider. CLI `--models mlx/*` only sets
- * the initial selector scope: Tab, `/models`, RPC enumeration, explicit model
- * resolution, and restored sessions still consult the registry's unscoped
- * reads. Filter those reads at their shared boundary before pi constructs its
- * runtime so every path sees only the exact local models this process serves.
+ * `mlx agent` is an offline/local product, but pi's runtime also composes every
+ * built-in cloud provider. CLI `--models mlx/*` only sets the initial selector
+ * scope: Tab, `/models`, RPC enumeration, explicit model resolution, and
+ * restored sessions all read the runtime's unscoped catalog/availability
+ * directly (the `ModelRegistry` facade handed to extensions delegates to the
+ * same runtime). Filter those reads at their shared boundary — the runtime
+ * prototype — so every path sees only the exact local models this process
+ * serves. Patching the runtime (not the extension-only facade) is what keeps
+ * the mlx-only guarantee across the selector / listing / resolution paths.
+ *
+ * The guarantee is an ALLOWLIST across three surfaces, all keyed on the `mlx`
+ * provider id:
+ *  1. Model reads (`getModels`/`getAvailable*`/`getModel`) — exact local-model
+ *     identity (`api === 'mlx' && baseUrl === 'mlx://local'`).
+ *  2. Provider/auth reads (`getProviders`/`getProvider`/`checkAuth`/`getAuth`/
+ *     `isUsingOAuth`/`hasConfiguredAuth`/`listCredentials`/`getProviderAuthStatus`)
+ *     — never surface, report configured, resolve auth for, or enumerate a
+ *     credential of any non-`mlx` provider. `getAuth` is the pivotal one: pi's
+ *     `/login`, `/logout`, and the built-in `/llama` command all resolve auth
+ *     through it, and its OAuth/`LLAMA_BASE_URL` `fetch` consults neither
+ *     `PI_OFFLINE` nor any allow-network flag; returning `undefined` for non-mlx
+ *     makes those commands fail before any network or second-model-host load.
+ *  3. Auth mutation / network (`login` rejected for non-mlx; `refresh` forced to
+ *     `allowNetwork: false` so an explicit `allowNetwork: true` — e.g. pi's
+ *     `update --models` package command — can never override offline mode).
+ *
+ * The `mlx` provider is registered with a literal apiKey and never needs
+ * `/login`; streaming still works because `prepareRequest` calls `getAuth(model)`
+ * with `model.provider === 'mlx'`, which passes through.
+ *
+ * Surfaces 1-3 patch the runtime's PUBLIC facade, but pi's own internals read the
+ * composed provider map `this.models` (pi-ai `ModelsImpl`) directly — a boundary
+ * the facade patches cannot reach (e.g. `refresh` resolving a command-backed
+ * cloud credential). So there is a 4th, structural surface: the `recomposeProvider`
+ * choke composes ONLY `mlx` into `this.models`, making every internal read
+ * mlx-only by construction. It requires installation before the runtime is
+ * constructed (which `runAgent` guarantees).
  *
  * Keep this adapter isolated: once pi exposes a first-class provider allowlist
  * in `MainOptions`, this file can be replaced by that option without touching
  * the provider or CLI layers.
  */
+import { MLX_API, MLX_API_KEY, MLX_BASE_URL, MLX_PROVIDER_ID } from './mlx-identity.js';
 const activePrototypes = new WeakSet();
+/**
+ * `getAuth` is intentionally NOT part of {@link FilterableModelRuntime}: pi types
+ * it with two overloads (`getAuth(providerId)` and `getAuth(model)`), which a
+ * single structural signature cannot capture without breaking the runtime→
+ * interface assignment. It is wrapped by name (mlx pins a fixed resolution; the
+ * original is never delegated to).
+ */
+/** Method names wrapped through the structural interface, plus the loose `getAuth`. */
+const STRUCTURAL_METHODS = [
+    'getModels',
+    'getAvailableSnapshot',
+    'getAvailable',
+    'getModel',
+    'getProvider',
+    'hasConfiguredAuth',
+    'checkAuth',
+    'isUsingOAuth',
+    'getProviders',
+    'listCredentials',
+    'getProviderAuthStatus',
+    'login',
+    'refresh',
+];
 function requireMethodDescriptor(prototype, name) {
     const descriptor = Object.getOwnPropertyDescriptor(prototype, name);
     if (!descriptor || typeof descriptor.value !== 'function' || descriptor.writable !== true) {
-        throw new Error(`mlx agent: incompatible pi ModelRegistry.${name}; expected a writable prototype method`);
+        throw new Error(`mlx agent: incompatible pi ModelRuntime.${name}; expected a writable prototype method`);
     }
     return descriptor;
 }
 /**
- * Install an exact local-model read policy for one `runAgent()` lifetime.
- * Returns an idempotent restore callback.
+ * Install an exact local-model / mlx-only-provider policy for one `runAgent()`
+ * lifetime. Returns an idempotent restore callback.
  */
-export function installMlxOnlyModelRegistryFilter(Registry, modelIds) {
-    const prototype = Registry.prototype;
+export function installMlxOnlyModelRegistryFilter(Runtime, modelIds) {
+    const prototype = Runtime.prototype;
     if (activePrototypes.has(prototype)) {
-        throw new Error('mlx agent: concurrent ModelRegistry filtering in one process is not supported');
+        throw new Error('mlx agent: concurrent ModelRuntime filtering in one process is not supported');
     }
     const allowedIds = new Set(modelIds);
-    const isAllowed = (model) => model.provider === 'mlx' && allowedIds.has(model.id) && model.api === 'mlx' && model.baseUrl === 'mlx://local';
-    const descriptors = {
-        getAll: requireMethodDescriptor(prototype, 'getAll'),
-        getAvailable: requireMethodDescriptor(prototype, 'getAvailable'),
-        find: requireMethodDescriptor(prototype, 'find'),
-        hasConfiguredAuth: requireMethodDescriptor(prototype, 'hasConfiguredAuth'),
-    };
-    const getAll = descriptors.getAll.value;
-    const getAvailable = descriptors.getAvailable.value;
-    const find = descriptors.find.value;
-    const hasConfiguredAuth = descriptors.hasConfiguredAuth.value;
+    const isAllowed = (model) => model.provider === MLX_PROVIDER_ID &&
+        allowedIds.has(model.id) &&
+        model.api === MLX_API &&
+        model.baseUrl === MLX_BASE_URL;
+    // Capture every original descriptor up front (fail-closed if pi's method shape
+    // changed), so restore can put them all back verbatim.
+    const originals = {};
+    for (const name of STRUCTURAL_METHODS)
+        originals[name] = requireMethodDescriptor(prototype, name);
+    originals.getAuth = requireMethodDescriptor(prototype, 'getAuth');
+    // `recomposeProvider` is a private funnel method; capture it by name (it is not
+    // part of the structural interface — see the choke wrapper below).
+    originals.recomposeProvider = requireMethodDescriptor(prototype, 'recomposeProvider');
+    const iface = originals;
+    const getModels = iface.getModels.value;
+    const getAvailableSnapshot = iface.getAvailableSnapshot.value;
+    const getAvailable = iface.getAvailable.value;
+    const getModel = iface.getModel.value;
+    const getProvider = iface.getProvider.value;
+    const hasConfiguredAuth = iface.hasConfiguredAuth.value;
+    const checkAuth = iface.checkAuth.value;
+    const isUsingOAuth = iface.isUsingOAuth.value;
+    const getProviders = iface.getProviders.value;
+    const listCredentials = iface.listCredentials.value;
+    const getProviderAuthStatus = iface.getProviderAuthStatus.value;
+    const login = iface.login.value;
+    const refresh = iface.refresh.value;
+    const recomposeProvider = originals.recomposeProvider.value;
     Object.defineProperties(prototype, {
-        getAll: {
-            ...descriptors.getAll,
+        getModels: {
+            ...originals.getModels,
+            value(providerId) {
+                return getModels.call(this, providerId).filter(isAllowed);
+            },
+        },
+        getAvailableSnapshot: {
+            ...originals.getAvailableSnapshot,
             value() {
-                return getAll.call(this).filter(isAllowed);
+                return getAvailableSnapshot.call(this).filter(isAllowed);
             },
         },
         getAvailable: {
-            ...descriptors.getAvailable,
-            value() {
-                return getAvailable.call(this).filter(isAllowed);
+            ...originals.getAvailable,
+            // The runtime read is async, so filter the resolved snapshot. Preserve the
+            // Promise contract (never turn a rejection into a filtered success).
+            value(providerId) {
+                return getAvailable.call(this, providerId).then((models) => models.filter(isAllowed));
             },
         },
-        find: {
-            ...descriptors.find,
+        getModel: {
+            ...originals.getModel,
             value(provider, modelId) {
-                if (provider !== 'mlx' || !allowedIds.has(modelId))
+                if (provider !== MLX_PROVIDER_ID || !allowedIds.has(modelId))
                     return undefined;
-                const model = find.call(this, provider, modelId);
+                const model = getModel.call(this, provider, modelId);
                 return model && isAllowed(model) ? model : undefined;
             },
         },
+        getProvider: {
+            ...originals.getProvider,
+            // Streaming uses `this.models.getProvider` (a different object); this only
+            // gates external reads (e.g. `/logout`). Non-mlx must never surface.
+            value(providerId) {
+                return providerId === MLX_PROVIDER_ID ? getProvider.call(this, providerId) : undefined;
+            },
+        },
         hasConfiguredAuth: {
-            ...descriptors.hasConfiguredAuth,
-            value(model) {
-                return isAllowed(model) && hasConfiguredAuth.call(this, model);
+            ...originals.hasConfiguredAuth,
+            // The runtime signature takes a providerId string (not a model), so gate on
+            // the provider id alone: only 'mlx' may ever report configured auth.
+            value(providerId) {
+                return providerId === MLX_PROVIDER_ID && hasConfiguredAuth.call(this, providerId);
+            },
+        },
+        checkAuth: {
+            ...originals.checkAuth,
+            value(providerId) {
+                return providerId === MLX_PROVIDER_ID ? checkAuth.call(this, providerId) : Promise.resolve(undefined);
+            },
+        },
+        isUsingOAuth: {
+            ...originals.isUsingOAuth,
+            value(providerId) {
+                return providerId === MLX_PROVIDER_ID && isUsingOAuth.call(this, providerId);
+            },
+        },
+        getProviders: {
+            ...originals.getProviders,
+            // `/login` enumerates from here; hide every cloud provider so only mlx is
+            // ever offered for sign-in.
+            value() {
+                return getProviders.call(this).filter((provider) => provider.id === MLX_PROVIDER_ID);
+            },
+        },
+        listCredentials: {
+            ...originals.listCredentials,
+            // `/logout` enumerates stored credentials from here (bypassing the composed
+            // model map); never reveal a non-mlx credential.
+            value() {
+                return listCredentials.call(this).then((creds) => creds.filter((cred) => cred.providerId === MLX_PROVIDER_ID));
+            },
+        },
+        getProviderAuthStatus: {
+            ...originals.getProviderAuthStatus,
+            // Reads the raw credential/config layer (not the model map); only mlx may
+            // report configured.
+            value(providerId) {
+                return providerId === MLX_PROVIDER_ID ? getProviderAuthStatus.call(this, providerId) : { configured: false };
+            },
+        },
+        getAuth: {
+            ...originals.getAuth,
+            // Pivotal offline gate + reserved-id invariant. `/login`, `/logout`, and the
+            // built-in `/llama` command resolve auth through the runtime's `getAuth`,
+            // whose OAuth / `LLAMA_BASE_URL` fetch ignores PI_OFFLINE — so non-mlx must
+            // resolve to nothing (string form `getAuth(id)` or model form via
+            // `model.provider`). For mlx we do NOT delegate to the composed provider:
+            // a persisted `models.json` `{ oauth:'radius', baseUrl }` overlay under the
+            // `mlx` id promotes a Radius builtin that merges a radius OAuth method into
+            // the composed `mlx` auth; with a stored/expired `mlx` oauth credential the
+            // real resolution would trigger an offline OAuth refresh throw ("OAuth
+            // refresh failed for mlx") and never reach our local stream. Pin mlx to the
+            // fixed local credential instead — identical to the no-overlay case, but
+            // immune to the overlay. prepareRequest reads only auth.apiKey/baseUrl(/headers)
+            // and dispatches streamSimple on the (api-matched) local closure.
+            value(providerOrModel, _overrides) {
+                const providerId = typeof providerOrModel === 'string'
+                    ? providerOrModel
+                    : providerOrModel?.provider;
+                if (providerId !== MLX_PROVIDER_ID)
+                    return Promise.resolve(undefined);
+                return Promise.resolve({ auth: { apiKey: MLX_API_KEY, baseUrl: MLX_BASE_URL } });
+            },
+        },
+        login: {
+            ...originals.login,
+            // Authoritative offline gate: reject any non-mlx login BEFORE dispatch, so
+            // the cloud OAuth `fetch` (e.g. radius.pi.dev) can never fire — even if some
+            // path passes a provider id directly, bypassing the filtered enumeration.
+            value(providerId, type, interaction) {
+                if (providerId !== MLX_PROVIDER_ID) {
+                    return Promise.reject(new Error('mlx agent is offline: provider login is disabled'));
+                }
+                return login.call(this, providerId, type, interaction);
+            },
+        },
+        refresh: {
+            ...originals.refresh,
+            // Force offline: pi's `refresh({ allowNetwork })` honours an explicit
+            // `true` (e.g. the `update --models` package command) even under
+            // PI_OFFLINE=1. Pin it off so no catalog fetch can escape the boundary.
+            value(options = {}) {
+                return refresh.call(this, { ...options, allowNetwork: false });
+            },
+        },
+        recomposeProvider: {
+            ...originals.recomposeProvider,
+            // CHOKE POINT (structural backbone). `recomposeProvider` is the single
+            // funnel that writes a provider into the runtime's model map `this.models`
+            // (pi's `rebuildProviders` sweep of builtins/config + the extension
+            // `registerProvider`/`registerNativeProvider`/`unregisterProvider` paths).
+            // Compose ONLY `mlx`, so `this.models` is mlx-only by construction. That
+            // closes the whole class of INTERNAL `this.models.*` reads the public gates
+            // above cannot reach — most importantly `refresh`'s availability pass, which
+            // resolves a configured cloud provider's credential (executing a
+            // command-backed apiKey) with no allowNetwork gate — plus the `/llama`
+            // getAuth and catalog-refresh paths. Requires install before the runtime is
+            // created (runAgent installs before pi.main constructs it); the mlx provider
+            // registers later via its own `recomposeProvider('mlx')`, allowed through.
+            value(providerId) {
+                if (providerId !== MLX_PROVIDER_ID)
+                    return;
+                recomposeProvider.call(this, providerId);
             },
         },
     });
@@ -75,7 +266,7 @@ export function installMlxOnlyModelRegistryFilter(Registry, modelIds) {
     return () => {
         if (restored)
             return;
-        Object.defineProperties(prototype, descriptors);
+        Object.defineProperties(prototype, originals);
         activePrototypes.delete(prototype);
         restored = true;
     };

package/dist/provider/models.d.ts

@@ -1,9 +1,8 @@
 /**
  * Local model discovery for the mlx pi provider.
  *
- * Ports the discovery walk from
- * `packages/cli/src/commands/launch-claude/discover.ts` (which is bin-only
- * and must not be imported from here; the cli copy stays untouched) and
+ * Ports the discovery walk from `@mlx-node/server/host`
+ * (`packages/server/src/host/discover.ts`; that copy stays untouched) and
  * pairs every discovered checkpoint with a pi `ProviderModelConfig` entry
  * ready for `pi.registerProvider('mlx', { models })`.
  *

package/dist/provider/models.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"models.d.ts","sourceRoot":"","sources":["../../src/provider/models.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;;;;;;;GAgBG;AAMH,OAAO,KAAK,EAAE,mBAAmB,EAAE,MAAM,iCAAiC,CAAC;AAG3E,OAAO,KAAK,EAAE,mBAAmB,EAAE,MAAM,aAAa,CAAC;AAGvD,6EAA6E;AAC7E,MAAM,WAAW,YAAY;IAC3B,UAAU,EAAE,mBAAmB,CAAC;IAChC,OAAO,EAAE,mBAAmB,CAAC;CAC9B;AAmED;;;;;;;;GAQG;AACH,wBAAsB,iBAAiB,CAAC,SAAS,EAAE,MAAM,GAAG,OAAO,CAAC,YAAY,EAAE,CAAC,CAsDlF"}
+{"version":3,"file":"models.d.ts","sourceRoot":"","sources":["../../src/provider/models.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;;;;;;GAeG;AAMH,OAAO,KAAK,EAAE,mBAAmB,EAAE,MAAM,iCAAiC,CAAC;AAG3E,OAAO,KAAK,EAAE,mBAAmB,EAAE,MAAM,aAAa,CAAC;AAGvD,6EAA6E;AAC7E,MAAM,WAAW,YAAY;IAC3B,UAAU,EAAE,mBAAmB,CAAC;IAChC,OAAO,EAAE,mBAAmB,CAAC;CAC9B;AAmHD;;;;;;;;GAQG;AACH,wBAAsB,iBAAiB,CAAC,SAAS,EAAE,MAAM,GAAG,OAAO,CAAC,YAAY,EAAE,CAAC,CAuDlF"}

package/dist/provider/models.js

@@ -1,9 +1,8 @@
 /**
  * Local model discovery for the mlx pi provider.
  *
- * Ports the discovery walk from
- * `packages/cli/src/commands/launch-claude/discover.ts` (which is bin-only
- * and must not be imported from here; the cli copy stays untouched) and
+ * Ports the discovery walk from `@mlx-node/server/host`
+ * (`packages/server/src/host/discover.ts`; that copy stays untouched) and
  * pairs every discovered checkpoint with a pi `ProviderModelConfig` entry
  * ready for `pi.registerProvider('mlx', { models })`.
  *
@@ -32,39 +31,65 @@ const FAMILY_TRAITS = {
     qwen3: { reasoning: true, fallbackContextWindow: 40960 },
     qwen3_5: { reasoning: true, fallbackContextWindow: 262144 },
     qwen3_5_moe: { reasoning: true, fallbackContextWindow: 262144 },
-    gemma4: { reasoning: false, fallbackContextWindow: 131072 },
+    gemma4: {
+        reasoning: true,
+        thinkingLevelMap: {
+            minimal: 'minimal',
+            low: null,
+            medium: null,
+            high: 'high',
+        },
+        fallbackContextWindow: 131072,
+    },
     lfm2: { reasoning: true, fallbackContextWindow: 128000 },
     lfm2_moe: { reasoning: true, fallbackContextWindow: 128000 },
 };
 function positiveInteger(value) {
     return typeof value === 'number' && Number.isFinite(value) && value > 0 ? Math.floor(value) : undefined;
 }
+function nonEmptyRecord(value) {
+    return typeof value === 'object' && value !== null && !Array.isArray(value) && Object.keys(value).length > 0;
+}
 /**
- * Read the trained context window from `<modelPath>/config.json`:
+ * Read cheap discovery metadata from `<modelPath>/config.json`.
+ *
+ * The trained context window comes from:
  * root `max_position_embeddings` first (qwen3, lfm2), then
  * `text_config.max_position_embeddings` (qwen3_5, qwen3_5_moe, gemma4
- * unified), else the family fallback. `detectModelType` already parsed
- * this file, so a read/parse failure here (e.g. a racing rewrite) lands
- * on the fallback instead of dropping the model.
+ * unified), else the family fallback.
+ *
+ * Image support is advertised only when a family with a native multimodal
+ * implementation carries its valid, non-empty vision marker: `vision_config`
+ * for Qwen, and either `vision_config` or `unified_vision_config` for Gemma.
+ * This lets Pi's model picker and `--list-models` expose checkpoint capability
+ * without loading weights. The first resident load remains authoritative and
+ * reconciles this optimistic config-level advertisement via
+ * `session.supportsImages()` (for example, when conversion stripped an
+ * incompatible vision tower).
+ *
+ * `detectModelType` already parsed this file, so a read/parse failure here
+ * (e.g. a racing rewrite) lands on the context fallback and text-only input
+ * instead of dropping the model or guessing a positive capability.
  */
-async function readContextWindow(modelPath, fallback) {
+async function readDiscoveryMetadata(modelPath, modelType, fallbackContextWindow) {
     try {
         const raw = await readFile(join(modelPath, 'config.json'), 'utf-8');
         const config = JSON.parse(raw);
         const root = positiveInteger(config.max_position_embeddings);
-        if (root !== undefined)
-            return root;
         const textConfig = config.text_config;
-        if (typeof textConfig === 'object' && textConfig !== null && !Array.isArray(textConfig)) {
-            const nested = positiveInteger(textConfig.max_position_embeddings);
-            if (nested !== undefined)
-                return nested;
-        }
+        const nested = nonEmptyRecord(textConfig) ? positiveInteger(textConfig.max_position_embeddings) : undefined;
+        const hasVisionConfig = nonEmptyRecord(config.vision_config);
+        const supportsImages = modelType === 'gemma4'
+            ? hasVisionConfig || nonEmptyRecord(config.unified_vision_config)
+            : (modelType === 'qwen3_5' || modelType === 'qwen3_5_moe') && hasVisionConfig;
+        return {
+            contextWindow: root ?? nested ?? fallbackContextWindow,
+            supportsImages,
+        };
     }
     catch {
-        // fall through to the family fallback
+        return { contextWindow: fallbackContextWindow, supportsImages: false };
     }
-    return fallback;
 }
 /**
  * Scan `modelsDir` for chat-capable model subdirectories and build their
@@ -113,14 +138,15 @@ export async function discoverMlxModels(modelsDir) {
             continue;
         }
         const name = basename(full);
-        const contextWindow = await readContextWindow(full, traits.fallbackContextWindow);
+        const { contextWindow, supportsImages } = await readDiscoveryMetadata(full, modelType, traits.fallbackContextWindow);
         out.push({
             discovered: { name, path: full, modelType },
             piModel: {
                 id: name,
                 name,
                 reasoning: traits.reasoning,
-                input: ['text'],
+                thinkingLevelMap: traits.thinkingLevelMap,
+                input: supportsImages ? ['text', 'image'] : ['text'],
                 cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
                 contextWindow,
                 maxTokens: preset.maxOutputTokens,

package/dist/provider/stream-adapter.d.ts

@@ -36,7 +36,7 @@
  *     layer is swallowed: there is no further recovery surface.
  */
 import type { Api, AssistantMessage, AssistantMessageEventStream, Context, Model, SimpleStreamOptions } from '@earendil-works/pi-ai';
-import type { ChatSession, PerformanceMetrics } from '@mlx-node/lm';
+import type { ChatSession, ChatStreamFinal, PerformanceMetrics } from '@mlx-node/lm';
 import type { DiscoveredModelLike } from '../types.js';
 /**
  * The exact `MlxModelHost` surface the adapter consumes, kept structural
@@ -46,8 +46,12 @@ import type { DiscoveredModelLike } from '../types.js';
 export interface StreamSimpleHost {
     /** Discovery record for `modelId` (source of the `ModelType` → launch preset). */
     modelInfo(modelId: string): DiscoveredModelLike | undefined;
-    /** Atomic resident selection + serialized inference closure (see `MlxModelHost`). */
-    runWithResident<T>(modelId: string, fn: (session: ChatSession) => Promise<T>): Promise<T>;
+    /**
+     * Atomic resident selection + serialized inference closure (see `MlxModelHost`).
+     * `fn` receives a `resident` boolean: `true` when the model was already warm
+     * (reused), `false` when this turn had to load/swap it.
+     */
+    runWithResident<T>(modelId: string, fn: (session: ChatSession, resident: boolean) => Promise<T>): Promise<T>;
     /** Flag the resident as post-error so the next turn does a full reset (see `MlxModelHost`). */
     markResidentDirty(modelId: string): void;
     /** Read-and-clear the resident's post-error flag; `true` ⇒ full-reset this turn. */
@@ -57,5 +61,34 @@ export interface StreamSimpleHost {
 }
 export type PerformanceRecorder = (message: AssistantMessage, performance: PerformanceMetrics) => void;
 export type RootCacheOwnerResolver = () => string | undefined;
-export declare function makeMlxStreamSimple(host: StreamSimpleHost, onPerformance?: PerformanceRecorder, resolveRootCacheOwner?: RootCacheOwnerResolver): (model: Model<Api>, context: Context, options?: SimpleStreamOptions) => AssistantMessageEventStream;
+/** Resolves the root session's JSONL path for the metrics-trace root correlation. */
+export type RootSessionFileResolver = () => string | undefined;
+/**
+ * Durable per-turn telemetry hook. Fires exactly once, only on a SUCCESSFUL
+ * native final (never on abort / error / load failure), from the one seam
+ * that sees the raw `ChatStreamFinal` alongside the per-request
+ * `options.sessionId` and the turn's minted `traceId`. Best-effort: the
+ * adapter guards the call so a throwing recorder can never break inference.
+ */
+export type TurnRecorder = (rec: {
+    traceId: string;
+    sessionId?: string;
+    /** Root session id/file snapshotted when this turn was submitted (see below). */
+    rootSessionId?: string;
+    rootSessionFile?: string;
+    model: string;
+    final: ChatStreamFinal;
+    /** Whole-turn wall-clock (ms): queue wait + resident selection + prefill + decode. */
+    durationMs: number;
+    /**
+     * Queue + cold-load wait (ms) BEFORE native work began this turn: the gap
+     * between turn submission and the serialized `runWithResident` callback
+     * firing (behind earlier inference and/or a model load/swap). Subtract from
+     * `durationMs` to recover execution-only time.
+     */
+    queueMs: number;
+    /** `true` when the model was already warm/resident, `false` on a cold load/swap this turn. */
+    resident: boolean;
+}) => void;
+export declare function makeMlxStreamSimple(host: StreamSimpleHost, onPerformance?: PerformanceRecorder, resolveRootCacheOwner?: RootCacheOwnerResolver, onTurnRecord?: TurnRecorder, onTurnStart?: () => void, resolveRootSessionFile?: RootSessionFileResolver): (model: Model<Api>, context: Context, options?: SimpleStreamOptions) => AssistantMessageEventStream;
 //# sourceMappingURL=stream-adapter.d.ts.map

package/dist/provider/stream-adapter.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"stream-adapter.d.ts","sourceRoot":"","sources":["../../src/provider/stream-adapter.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;GAoCG;AAEH,OAAO,KAAK,EACV,GAAG,EACH,gBAAgB,EAChB,2BAA2B,EAC3B,OAAO,EACP,KAAK,EACL,mBAAmB,EACpB,MAAM,uBAAuB,CAAC;AAE/B,OAAO,KAAK,EAAE,WAAW,EAAE,kBAAkB,EAAE,MAAM,cAAc,CAAC;AAEpE,OAAO,KAAK,EAAE,mBAAmB,EAAE,MAAM,aAAa,CAAC;AAOvD;;;;GAIG;AACH,MAAM,WAAW,gBAAgB;IAC/B,kFAAkF;IAClF,SAAS,CAAC,OAAO,EAAE,MAAM,GAAG,mBAAmB,GAAG,SAAS,CAAC;IAC5D,qFAAqF;IACrF,eAAe,CAAC,CAAC,EAAE,OAAO,EAAE,MAAM,EAAE,EAAE,EAAE,CAAC,OAAO,EAAE,WAAW,KAAK,OAAO,CAAC,CAAC,CAAC,GAAG,OAAO,CAAC,CAAC,CAAC,CAAC;IAC1F,+FAA+F;IAC/F,iBAAiB,CAAC,OAAO,EAAE,MAAM,GAAG,IAAI,CAAC;IACzC,oFAAoF;IACpF,oBAAoB,CAAC,OAAO,EAAE,MAAM,GAAG,OAAO,CAAC;IAC/C,gFAAgF;IAChF,kBAAkB,CAAC,OAAO,EAAE,MAAM,GAAG,IAAI,CAAC;CAC3C;AAED,MAAM,MAAM,mBAAmB,GAAG,CAAC,OAAO,EAAE,gBAAgB,EAAE,WAAW,EAAE,kBAAkB,KAAK,IAAI,CAAC;AACvG,MAAM,MAAM,sBAAsB,GAAG,MAAM,MAAM,GAAG,SAAS,CAAC;AAqF9D,wBAAgB,mBAAmB,CACjC,IAAI,EAAE,gBAAgB,EACtB,aAAa,CAAC,EAAE,mBAAmB,EACnC,qBAAqB,CAAC,EAAE,sBAAsB,GAC7C,CAAC,KAAK,EAAE,KAAK,CAAC,GAAG,CAAC,EAAE,OAAO,EAAE,OAAO,EAAE,OAAO,CAAC,EAAE,mBAAmB,KAAK,2BAA2B,CAkOrG"}
+{"version":3,"file":"stream-adapter.d.ts","sourceRoot":"","sources":["../../src/provider/stream-adapter.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;GAoCG;AAEH,OAAO,KAAK,EACV,GAAG,EACH,gBAAgB,EAChB,2BAA2B,EAC3B,OAAO,EACP,KAAK,EACL,mBAAmB,EACpB,MAAM,uBAAuB,CAAC;AAE/B,OAAO,KAAK,EAAE,WAAW,EAAE,eAAe,EAAE,kBAAkB,EAAE,MAAM,cAAc,CAAC;AAErF,OAAO,KAAK,EAAE,mBAAmB,EAAE,MAAM,aAAa,CAAC;AAOvD;;;;GAIG;AACH,MAAM,WAAW,gBAAgB;IAC/B,kFAAkF;IAClF,SAAS,CAAC,OAAO,EAAE,MAAM,GAAG,mBAAmB,GAAG,SAAS,CAAC;IAC5D;;;;OAIG;IACH,eAAe,CAAC,CAAC,EAAE,OAAO,EAAE,MAAM,EAAE,EAAE,EAAE,CAAC,OAAO,EAAE,WAAW,EAAE,QAAQ,EAAE,OAAO,KAAK,OAAO,CAAC,CAAC,CAAC,GAAG,OAAO,CAAC,CAAC,CAAC,CAAC;IAC7G,+FAA+F;IAC/F,iBAAiB,CAAC,OAAO,EAAE,MAAM,GAAG,IAAI,CAAC;IACzC,oFAAoF;IACpF,oBAAoB,CAAC,OAAO,EAAE,MAAM,GAAG,OAAO,CAAC;IAC/C,gFAAgF;IAChF,kBAAkB,CAAC,OAAO,EAAE,MAAM,GAAG,IAAI,CAAC;CAC3C;AAED,MAAM,MAAM,mBAAmB,GAAG,CAAC,OAAO,EAAE,gBAAgB,EAAE,WAAW,EAAE,kBAAkB,KAAK,IAAI,CAAC;AACvG,MAAM,MAAM,sBAAsB,GAAG,MAAM,MAAM,GAAG,SAAS,CAAC;AAC9D,qFAAqF;AACrF,MAAM,MAAM,uBAAuB,GAAG,MAAM,MAAM,GAAG,SAAS,CAAC;AAE/D;;;;;;GAMG;AACH,MAAM,MAAM,YAAY,GAAG,CAAC,GAAG,EAAE;IAC/B,OAAO,EAAE,MAAM,CAAC;IAChB,SAAS,CAAC,EAAE,MAAM,CAAC;IACnB,iFAAiF;IACjF,aAAa,CAAC,EAAE,MAAM,CAAC;IACvB,eAAe,CAAC,EAAE,MAAM,CAAC;IACzB,KAAK,EAAE,MAAM,CAAC;IACd,KAAK,EAAE,eAAe,CAAC;IACvB,sFAAsF;IACtF,UAAU,EAAE,MAAM,CAAC;IACnB;;;;;OAKG;IACH,OAAO,EAAE,MAAM,CAAC;IAChB,8FAA8F;IAC9F,QAAQ,EAAE,OAAO,CAAC;CACnB,KAAK,IAAI,CAAC;AAqFX,wBAAgB,mBAAmB,CACjC,IAAI,EAAE,gBAAgB,EACtB,aAAa,CAAC,EAAE,mBAAmB,EACnC,qBAAqB,CAAC,EAAE,sBAAsB,EAC9C,YAAY,CAAC,EAAE,YAAY,EAC3B,WAAW,CAAC,EAAE,MAAM,IAAI,EACxB,sBAAsB,CAAC,EAAE,uBAAuB,GAC/C,CAAC,KAAK,EAAE,KAAK,CAAC,GAAG,CAAC,EAAE,OAAO,EAAE,OAAO,EAAE,OAAO,CAAC,EAAE,mBAAmB,KAAK,2BAA2B,CAuRrG"}

package/dist/provider/stream-adapter.js

@@ -125,7 +125,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 +138,7 @@ export function makeMlxStreamSimple(host, onPerformance, resolveRootCacheOwner)
         let emitter;
         let signal;
         let rootCacheOwnerId;
+        let rootSessionFile;
         let resolvedReasoning;
         let detachAbort;
         /**
@@ -208,8 +209,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,12 +242,36 @@ 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;
+                // 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.
+                    }
+                }
                 publishEffectiveContextWindow(model, session);
                 const supportsImages = publishImageCapability(model, session);
                 const discovered = host.modelInfo(model.id);
@@ -308,6 +336,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));