@mlx-node/agent 0.0.7

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

@@ -0,0 +1,73 @@
+/**
+ * Agent-private `ChatSession` warm-reuse helper.
+ *
+ * Port of `packages/server/src/chat-session-warm-reuse.ts` — that module
+ * is deliberately kept off the server's export map, so the agent package
+ * carries its own copy with identical runtime behavior. Keep the two in
+ * sync.
+ *
+ * Why this helper exists at all: `ChatSession.reset()` is the safe
+ * public wipe — it always calls `model.resetCaches()` because the
+ * underlying `SessionCapableModel` may be shared across session
+ * lifetimes. The agent provider bridge, however, owns exactly one
+ * session per model process and replays pi's full message history on
+ * every LLM call, so the native KV cache always belongs to the chain
+ * being replayed. A JS-state-only reset that preserves the native
+ * cache is correct there: the next `primeHistory()` +
+ * `startFromHistoryStream()` lets the native prefix verifier recover
+ * 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
+ * `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';
+/**
+ * Private structural view of the `ChatSession` JS-side state that the
+ * warm-reuse helper needs to wipe. Mirrors the internal state
+ * documented on `ChatSession` itself — field names are load-bearing:
+ * they must byte-match the concrete class's private fields or the
+ * cast-based mutation below silently no-ops.
+ */
+interface ChatSessionWarmReuseInternals {
+    inFlight: boolean;
+    history: unknown[];
+    lastImagesKey: string | null;
+    lastAudioKey: string | null;
+    turnCount: number;
+    unresolvedOkToolCallCount: number | null;
+    needsFullReplay: boolean;
+}
+/**
+ * Runtime list of the `ChatSession` private field names this module
+ * reads or writes — exported solely so the drift test can assert each
+ * one still exists on a real `ChatSession` instance.
+ */
+export declare const WARM_REUSE_TOUCHED_FIELDS: ReadonlyArray<keyof ChatSessionWarmReuseInternals>;
+/**
+ * JS-state-only reset that DELIBERATELY preserves the underlying
+ * model's native KV cache and `cached_token_history`.
+ *
+ * @internal agent-private — used only by the provider bridge's
+ * per-call warm replay (`resetPreservingNativeCacheForWarmReuse` →
+ * `primeHistory` → `startFromHistoryStream`). Never export from this
+ * package's `index.ts`.
+ *
+ * Wipes ONLY the JS-side session state (history array, image key, turn
+ * counter, tool-call fan-out guard). With this function, the JS session
+ * is fresh enough for `ChatSession.primeHistory()` (which requires
+ * `turnCount === 0`) while the native prefix verifier can still recover
+ * the reused prefix on the next `chatSessionStart` and skip the
+ * corresponding re-prefill.
+ */
+export declare function resetPreservingNativeCacheForWarmReuse<M extends SessionCapableModel>(session: ChatSession<M>): Promise<void>;
+export {};
+//# sourceMappingURL=warm-reuse.d.ts.map

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

@@ -0,0 +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"}

package/dist/provider/warm-reuse.js

@@ -0,0 +1,88 @@
+/**
+ * Agent-private `ChatSession` warm-reuse helper.
+ *
+ * Port of `packages/server/src/chat-session-warm-reuse.ts` — that module
+ * is deliberately kept off the server's export map, so the agent package
+ * carries its own copy with identical runtime behavior. Keep the two in
+ * sync.
+ *
+ * Why this helper exists at all: `ChatSession.reset()` is the safe
+ * public wipe — it always calls `model.resetCaches()` because the
+ * underlying `SessionCapableModel` may be shared across session
+ * lifetimes. The agent provider bridge, however, owns exactly one
+ * session per model process and replays pi's full message history on
+ * every LLM call, so the native KV cache always belongs to the chain
+ * being replayed. A JS-state-only reset that preserves the native
+ * cache is correct there: the next `primeHistory()` +
+ * `startFromHistoryStream()` lets the native prefix verifier recover
+ * 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
+ * `packages/agent/__test__/warm-reuse.test.ts` checks every name in
+ * {@link WARM_REUSE_TOUCHED_FIELDS} against a real `ChatSession`
+ * instance.
+ */
+/**
+ * `Record<keyof ..., true>` forces this map to list EXACTLY the fields of
+ * {@link ChatSessionWarmReuseInternals}: adding/removing/renaming a field
+ * in the interface without updating the map is a compile error, so the
+ * runtime list below can never drift from the fields the helper touches.
+ */
+const WARM_REUSE_TOUCHED_FIELD_SET = {
+    inFlight: true,
+    history: true,
+    lastImagesKey: true,
+    lastAudioKey: true,
+    turnCount: true,
+    unresolvedOkToolCallCount: true,
+    needsFullReplay: true,
+};
+/**
+ * Runtime list of the `ChatSession` private field names this module
+ * reads or writes — exported solely so the drift test can assert each
+ * one still exists on a real `ChatSession` instance.
+ */
+export const WARM_REUSE_TOUCHED_FIELDS = Object.keys(WARM_REUSE_TOUCHED_FIELD_SET);
+/**
+ * JS-state-only reset that DELIBERATELY preserves the underlying
+ * model's native KV cache and `cached_token_history`.
+ *
+ * @internal agent-private — used only by the provider bridge's
+ * per-call warm replay (`resetPreservingNativeCacheForWarmReuse` →
+ * `primeHistory` → `startFromHistoryStream`). Never export from this
+ * package's `index.ts`.
+ *
+ * Wipes ONLY the JS-side session state (history array, image key, turn
+ * counter, tool-call fan-out guard). With this function, the JS session
+ * is fresh enough for `ChatSession.primeHistory()` (which requires
+ * `turnCount === 0`) while the native prefix verifier can still recover
+ * the reused prefix on the next `chatSessionStart` and skip the
+ * corresponding re-prefill.
+ */
+export async function resetPreservingNativeCacheForWarmReuse(session) {
+    // TypeScript `private` fields are only compile-time checks; at
+    // runtime they are ordinary properties. The cast through
+    // `ChatSessionWarmReuseInternals` preserves full static typing for
+    // this helper's mutations while bypassing the `private` gate — which
+    // is correct here because this helper is the designated agent-side
+    // friend accessor. The cast is funneled through `unknown` because TS
+    // correctly rejects a direct `ChatSession → Internals` cast when the
+    // concrete class has other non-internals fields.
+    const internals = session;
+    if (internals.inFlight) {
+        throw new Error('ChatSession: cannot resetPreservingNativeCacheForWarmReuse() while a send() is in flight; await the previous call first');
+    }
+    internals.history = [];
+    internals.lastImagesKey = null;
+    internals.lastAudioKey = null;
+    internals.turnCount = 0;
+    internals.unresolvedOkToolCallCount = null;
+    internals.needsFullReplay = false;
+}

package/dist/run-agent.d.ts

@@ -0,0 +1,62 @@
+/**
+ * `runAgent` — the boot shell that hands control to pi's `main()` with
+ * the mlx provider, permission gate, local subagents, and terminal branding installed.
+ *
+ * Spike-proven boot contract:
+ * - The env vars below must be set BEFORE any runtime import of
+ *   `@earendil-works/pi-coding-agent` (pi reads its config env at
+ *   import/call time) — hence the dynamic import and the type-only
+ *   top-level pi import here.
+ * - pi's `main()` RETURNS on the happy path but `process.exit()`s on
+ *   help/error/package-command paths, so nothing critical may run after
+ *   `await main()`.
+ * - In print/json mode pi takes over stdout (our writes are rerouted to
+ *   stderr) and reads non-TTY stdin to EOF as prompt input — this
+ *   wrapper must never consume or hold stdin. Paged-config temp cleanup in
+ *   `finally` is best-effort for normal returns/rejections, never correctness-
+ *   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 { MlxModelInfo } from './provider/models.js';
+/** Shape of pi's `main(argv, { extensionFactories })` — also the test seam. */
+export type RunAgentMain = (args: string[], opts: {
+    extensionFactories: InlineExtension[];
+}) => Promise<void>;
+export interface RunAgentPi {
+    main: RunAgentMain;
+    ModelRegistry: FilterableModelRegistryConstructor;
+}
+/** @internal Narrow lifecycle seam for the agent's paged config overlays. */
+export interface AgentPagedConfigOverrides {
+    resolve(modelPath: string, modelType?: string): Promise<string>;
+    cleanup(): Promise<void>;
+}
+export interface RunAgentOptions {
+    /** Resolved models directory (context for callers/diagnostics — discovery already ran). */
+    modelsDir: string;
+    /** Discovered models to serve through the in-process `mlx` provider. */
+    models: MlxModelInfo[];
+    /** Passthrough args handed to pi's `main()` verbatim. */
+    argv: string[];
+    /** Native inference-log path to surface after Pi takes over the TUI. */
+    traceLogFile?: string;
+    /** @internal Test seam; when set, the pi dynamic import is skipped entirely. */
+    piImpl?: RunAgentPi;
+    /** @internal Test seam for paged model-path resolution and cleanup. */
+    pagedConfigOverrides?: AgentPagedConfigOverrides;
+}
+/** @internal Exact opt-in parser kept separate so non-`1` values stay disabled. */
+export declare function agentGemmaDraftEnabled(env?: NodeJS.ProcessEnv): boolean;
+/**
+ * Seed the pi/mlx environment (never clobbering user-set values) and run
+ * pi's `main()` with the mlx inline extensions. May not return: pi
+ * `process.exit()`s on help/error paths.
+ *
+ * This is a dedicated-process entrypoint: pi owns stdin/stdout, signal
+ * handlers, and process exit while `main()` runs. Do not run it concurrently
+ * with another pi SDK/runtime in the same process. The temporary registry
+ * policy is restored when `main()` returns or rejects.
+ */
+export declare function runAgent(opts: RunAgentOptions): Promise<void>;
+//# sourceMappingURL=run-agent.d.ts.map

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

@@ -0,0 +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"}

package/dist/run-agent.js

@@ -0,0 +1,86 @@
+/**
+ * `runAgent` — the boot shell that hands control to pi's `main()` with
+ * the mlx provider, permission gate, local subagents, and terminal branding installed.
+ *
+ * Spike-proven boot contract:
+ * - The env vars below must be set BEFORE any runtime import of
+ *   `@earendil-works/pi-coding-agent` (pi reads its config env at
+ *   import/call time) — hence the dynamic import and the type-only
+ *   top-level pi import here.
+ * - pi's `main()` RETURNS on the happy path but `process.exit()`s on
+ *   help/error/package-command paths, so nothing critical may run after
+ *   `await main()`.
+ * - In print/json mode pi takes over stdout (our writes are rerouted to
+ *   stderr) and reads non-TTY stdin to EOF as prompt input — this
+ *   wrapper must never consume or hold stdin. Paged-config temp cleanup in
+ *   `finally` is best-effort for normal returns/rejections, never correctness-
+ *   critical on pi's hard `process.exit()` paths.
+ */
+import { homedir } from 'node:os';
+import { join } from 'node:path';
+import { PagedConfigOverrideManager } from '@mlx-node/lm';
+import { createPermissionGateExtension } from './extensions/permission-gate.js';
+import { createSubagentExtension } from './extensions/subagent.js';
+import { createTerminalTitleExtension } from './extensions/terminal-title.js';
+import { createTraceNoticeExtension } from './extensions/trace-notice.js';
+import { createMlxProviderExtension } from './provider/index.js';
+import { MlxModelHost } from './provider/model-host.js';
+import { installMlxOnlyModelRegistryFilter, } from './provider/model-registry-filter.js';
+/** @internal Exact opt-in parser kept separate so non-`1` values stay disabled. */
+export function agentGemmaDraftEnabled(env = process.env) {
+    return env.MLX_AGENT_ENABLE_GEMMA_DRAFT === '1';
+}
+/**
+ * Seed the pi/mlx environment (never clobbering user-set values) and run
+ * pi's `main()` with the mlx inline extensions. May not return: pi
+ * `process.exit()`s on help/error paths.
+ *
+ * This is a dedicated-process entrypoint: pi owns stdin/stdout, signal
+ * handlers, and process exit while `main()` runs. Do not run it concurrently
+ * with another pi SDK/runtime in the same process. The temporary registry
+ * policy is restored when `main()` returns or rejects.
+ */
+export async function runAgent(opts) {
+    process.env.PI_CODING_AGENT_DIR ??= join(homedir(), '.mlx-node', 'agent');
+    process.env.PI_SKIP_VERSION_CHECK ??= '1';
+    // 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';
+    // 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
+    // paged overlay intentionally hides an embedded draft/ directory: the
+    // current speculative executor is flat-cache-only and can regress quantized
+    // agent workloads. Users may explicitly opt back into that native behavior.
+    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),
+        requirePagedCache: true,
+    });
+    // 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.
+    const pi = opts.piImpl ?? (await import('@earendil-works/pi-coding-agent'));
+    const restoreModelRegistry = installMlxOnlyModelRegistryFilter(pi.ModelRegistry, 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),
+                createPermissionGateExtension(),
+                ...(subagentsEnabled ? [createSubagentExtension()] : []),
+                ...(opts.traceLogFile !== undefined ? [createTraceNoticeExtension(opts.traceLogFile)] : []),
+                createTerminalTitleExtension(),
+            ],
+        });
+    }
+    finally {
+        try {
+            restoreModelRegistry();
+        }
+        finally {
+            await pagedConfigOverrides.cleanup();
+        }
+    }
+}

package/dist/types.d.ts

@@ -0,0 +1,8 @@
+import type { ModelType } from '@mlx-node/lm';
+/** Structural mirror of the CLI's DiscoveredModel — avoids a cli↔agent dependency cycle. */
+export interface DiscoveredModelLike {
+    name: string;
+    path: string;
+    modelType: ModelType;
+}
+//# sourceMappingURL=types.d.ts.map

package/dist/types.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"types.d.ts","sourceRoot":"","sources":["../src/types.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,SAAS,EAAE,MAAM,cAAc,CAAC;AAE9C,4FAA4F;AAC5F,MAAM,WAAW,mBAAmB;IAClC,IAAI,EAAE,MAAM,CAAC;IACb,IAAI,EAAE,MAAM,CAAC;IACb,SAAS,EAAE,SAAS,CAAC;CACtB"}

package/dist/types.js

@@ -0,0 +1 @@
+export {};

package/package.json

@@ -0,0 +1,42 @@
+{
+  "name": "@mlx-node/agent",
+  "version": "0.0.7",
+  "homepage": "https://github.com/mlx-node/mlx-node",
+  "bugs": {
+    "url": "https://github.com/mlx-node/mlx-node/issues"
+  },
+  "license": "MIT",
+  "repository": {
+    "type": "git",
+    "url": "https://github.com/mlx-node/mlx-node.git",
+    "directory": "packages/agent"
+  },
+  "files": [
+    "dist"
+  ],
+  "type": "module",
+  "main": "./dist/index.js",
+  "types": "./dist/index.d.ts",
+  "exports": {
+    ".": {
+      "types": "./dist/index.d.ts",
+      "import": "./dist/index.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.7",
+    "@mlx-node/server": "0.0.7",
+    "typebox": "1.3.6"
+  },
+  "devDependencies": {
+    "@types/node": "^26.0.0"
+  },
+  "engines": {
+    "node": ">=22.19.0"
+  }
+}