@mlx-node/agent 0.0.12 → 0.0.15

package/src/provider/warm-reuse.ts

@@ -0,0 +1,125 @@
+/**
+ * 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`, `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 { 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
+ * 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;
+  defaultConfig?: ChatConfig;
+  activeTools: ChatConfig['tools'];
+}
+
+/**
+ * `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: Record<keyof ChatSessionWarmReuseInternals, true> = {
+  inFlight: true,
+  history: true,
+  lastImagesKey: true,
+  lastAudioKey: true,
+  turnCount: true,
+  unresolvedOkToolCallCount: true,
+  needsFullReplay: true,
+  defaultConfig: true,
+  activeTools: 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) as 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 async function resetPreservingNativeCacheForWarmReuse<M extends SessionCapableModel>(
+  session: ChatSession<M>,
+): Promise<void> {
+  // 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 as unknown as ChatSessionWarmReuseInternals;
+  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;
+  // 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/src/run-agent.ts

@@ -0,0 +1,178 @@
+/**
+ * `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 type { InlineExtension } from '@earendil-works/pi-coding-agent';
+import { coldCacheDrain } from '@mlx-node/core';
+import { PagedConfigOverrideManager } from '@mlx-node/lm';
+
+import { createDelegationExtension } from './extensions/delegation.js';
+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';
+import { createTraceNoticeExtension } from './extensions/trace-notice.js';
+import { createMlxProviderExtension } from './provider/index.js';
+import { MlxModelHost } from './provider/model-host.js';
+import {
+  type FilterableModelRuntimeConstructor,
+  installMlxOnlyModelRegistryFilter,
+} 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;
+  ModelRuntime: FilterableModelRuntimeConstructor;
+}
+
+/** @internal Narrow lifecycle seam for the agent's paged config overlays. */
+export interface AgentPagedConfigOverrides {
+  resolve(modelPath: string, modelType?: string, persistPagedCache?: boolean): Promise<string>;
+  cleanup(): Promise<void>;
+}
+
+export interface RunAgentOptions {
+  /** Focused worker tools/permissions; inference and session defaults remain shared. */
+  mode?: 'delegate';
+  /** Explicit approval supplied by the outer caller, captured before the worker starts. */
+  delegateCallerApproved?: boolean;
+  /** 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;
+  /**
+   * 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. */
+  pagedConfigOverrides?: AgentPagedConfigOverrides;
+}
+
+/** @internal Exact opt-in parser kept separate so non-`1` values stay disabled. */
+export function agentGemmaDraftEnabled(env: NodeJS.ProcessEnv = process.env): boolean {
+  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: RunAgentOptions): Promise<void> {
+  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';
+  // 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
+  // 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, 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
+  // 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: RunAgentPi = opts.piImpl ?? (await import('@earendil-works/pi-coding-agent'));
+  const restoreModelRegistry = installMlxOnlyModelRegistryFilter(
+    pi.ModelRuntime,
+    opts.models.map((model) => model.discovered.name),
+  );
+  const subagentsEnabled =
+    opts.mode !== 'delegate' &&
+    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(),
+        opts.mode === 'delegate'
+          ? createDelegationExtension({ callerApproved: opts.delegateCallerApproved })
+          : createPermissionGateExtension(),
+        ...(subagentsEnabled ? [createSubagentExtension()] : []),
+        ...(opts.traceLogFile !== undefined ? [createTraceNoticeExtension(opts.traceLogFile)] : []),
+        createTerminalTitleExtension(),
+      ],
+    });
+  } finally {
+    try {
+      restoreModelRegistry();
+    } 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/src/types.ts

@@ -0,0 +1,10 @@
+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;
+  /** External speculative drafter. When supplied, forwarded unchanged for loader validation. */
+  draftModelPath?: string;
+}