@bendyline/gezel-inference 0.0.0-stage → 1.0.0

package/dist/providers/mlx/provider.d.ts

@@ -0,0 +1,404 @@
+/**
+ * MlxProvider — talks to a local `mlx_lm.server` subprocess over HTTP.
+ * The server speaks OpenAI-compatible `/v1/chat/completions` with SSE
+ * streaming; tool calling rides the standard OpenAI function-tool
+ * shape, same as llama-cpp.
+ *
+ * Platform: **darwin-arm64 only**. Apple Silicon's unified memory +
+ * Metal Performance Shaders are what make MLX fast; there's no MLX
+ * backend for Intel Mac, Linux, or Windows. `ChatManager.ensureProvider`
+ * gates the branch on platform; this file assumes Apple Silicon.
+ *
+ * Session model is stateless (same as llama-cpp / Ollama): full
+ * transcript lives in memory and is re-sent per turn.
+ * `ProviderSessionState` is empty — the session record on disk is the
+ * source of truth.
+ *
+ * Lifecycle of the mlx_lm.server process is managed by a
+ * {@link NativeEngineSupervisor}: lazy-start on first turn, idle-stop,
+ * health-watch, restart budget. Startup is usually slower than
+ * llama-cpp (Python import + PyTorch/MLX init) so we give it a 180s
+ * startup budget.
+ *
+ * First-time install of mlx-lm itself is handled outside this file —
+ * `buildMlxProvider` in ChatManager calls `UvRuntime.ensureVenv('mlx',
+ * [mlxLmSpec])` before spawning the supervisor, ensuring the venv +
+ * packages exist.
+ */
+import type { ResolvedModelProfile } from '../../model-profile/types.js';
+import type { NativeEngineLifecycleSnapshot, NativeEngineSupervisor } from '../native/supervisor.js';
+import { ProviderQueue } from '../queue.js';
+import { type EnginePhaseEvent, type EngineStatsEvent, StreamingSessionBase } from '../streaming-session.js';
+import { type ToolExecutor } from '../tool-executor.js';
+import type { BatchCapability, ExternalToolCall, ExternalToolSpec, LLMProvider, LLMSession, ModelInfo, ProviderSessionState, SendAndWaitOpts, SessionOpts, ToolSurfaceSize, WireTranscriptEntry } from '../types.js';
+import { isSuccessfulAsyncHandoff } from './async-file-handoff.js';
+export { isSuccessfulAsyncHandoff };
+import { type EngineRequestPriority } from './engine-gate.js';
+import { type TemplateReasoningOpen } from './reasoning-stream.js';
+import { buildMidStreamDropMessage, computePreFirstByteBudgetMs, estimatePromptTokens, isMidStreamConnectionDrop } from './runtime-diagnostics.js';
+import { type MlxFatalError } from './stdout-parser.js';
+import { MlxVisionMode } from './vision-mode.js';
+export { buildMidStreamDropMessage, computePreFirstByteBudgetMs, estimatePromptTokens, isMidStreamConnectionDrop, };
+export declare class MlxProvider implements LLMProvider {
+    readonly name: "mlx";
+    readonly visionMode: MlxVisionMode;
+    readonly queue: ProviderQueue;
+    readonly supportsExternalTools = true;
+    readonly supportsPriorMessages = true;
+    private readonly supervisor?;
+    private readonly externalBaseUrl?;
+    private readonly defaultModel;
+    private readonly numCtx;
+    private readonly plannedReservation?;
+    private readonly fetchImpl;
+    /** Set in {@link shutdown}; a disposed provider must never respawn its engine. */
+    private disposed;
+    /**
+     * Catalog manager — when set, `listModels()` enumerates every model
+     * the user has installed via Settings → This Mac so the model picker
+     * in Settings → AI can pick from the full list. Optional because
+     * external-baseUrl mode (`GEZEL_MLX_SERVER_URL`) skips the catalog
+     * entirely. The picker selection writes back to
+     * `config.defaultModel.mlx` as a catalog id; `buildMlxProvider`
+     * resolves the id to a `modelDir` on the next provider rebuild.
+     */
+    private readonly modelManager?;
+    /**
+     * Friendly model name for the engine pill — e.g. "Qwen 3.5 8B".
+     * Surfaced in `loading_model` phase events as "Loading model
+     * (<name>)" so the user knows *which* MLX model is loading. Empty
+     * when the catalog manager wasn't wired in (external base URL mode)
+     * or when the runtime path doesn't map to a catalog entry — the
+     * pill falls back to the bare "Loading model" label in those cases.
+     */
+    private readonly modelDisplayName?;
+    /**
+     * Catalog model id for the on-disk model the provider is serving —
+     * `qwen3.6-27b-mlx`, `gemma4-e4b-mlx`, etc. Distinct from
+     * `defaultModel`, which holds the absolute model directory path
+     * (mlx_vlm.server's PRELOAD_MODEL value). Surfaced via
+     * {@link getEffectiveModelId} so the chat manager can resolve a
+     * tier for sessions that don't have an explicit model selection
+     * (`record.model` undefined AND `config.defaultModel.mlx`
+     * undefined). Without this, those sessions classify as
+     * `tier:tiny` because the tier resolver gets a path it can't parse.
+     *
+     * Empty when the provider was constructed in external-baseUrl mode
+     * (we don't know what the user's mlx-server is serving) or before
+     * the catalog has resolved a default model.
+     */
+    private readonly catalogModelId?;
+    /** See the constructor option of the same name. */
+    private readonly templateOpensReasoning?;
+    private readonly activeSessions;
+    /**
+     * Engine stdout → per-session phase events, plus the "is this line fatal?"
+     * state machine. See {@link EngineLogRouter}; the rules there are load-bearing.
+     */
+    private readonly logRouter;
+    private readonly engineStats;
+    /** Width-N gate over physical engine requests; see {@link MlxEngineGate}. */
+    private readonly engineGate;
+    private readonly batchMaxConcurrency;
+    private readonly turnProtection;
+    constructor(opts: {
+        /** A fixed tower: true only when the engine behind `baseUrl` serves images. */
+        visionEnabled?: boolean;
+        /** Supervised engines: text-only until a request carries images; see vision-mode.ts. */
+        vision?: MlxVisionMode;
+        supervisor?: NativeEngineSupervisor;
+        baseUrl?: string;
+        defaultModel?: string;
+        numCtx?: number;
+        /**
+         * Broker-ledger reservation for this replica: resident weights plus
+         * KV at the admitted window, computed by the build-time admission
+         * pass (see `LLMProvider.plannedReservationBytes`).
+         */
+        plannedReservationBytes?: number;
+        concurrency?: number;
+        /**
+         * Engine batch width — how many sequences the wrapped MLX server can
+         * generate concurrently. This defaults to 1 (a singleton BatchGenerator
+         * queue); raising it widens the gate and switches the queue to the
+         * adaptive interactive policy.
+         */
+        batchMaxConcurrency?: number;
+        affinity?: boolean;
+        fetchImpl?: typeof fetch;
+        modelManager?: import('./models.js').MlxModelManager;
+        modelDisplayName?: string;
+        /** See {@link MlxProvider.catalogModelId} for the contract. */
+        catalogModelId?: string;
+        /**
+         * Whether this model's chat template opens a reasoning block as the last
+         * thing it emits, and whether `enable_thinking` controls it. Detected once
+         * from the model directory at build time — see {@link templateOpensReasoning}.
+         */
+        templateOpensReasoning?: TemplateReasoningOpen;
+    });
+    /**
+     * Batch capability advertised by the wrapped MLX BatchGenerator.
+     * See {@link LLMProvider.batch}.
+     */
+    get batch(): BatchCapability;
+    /** Can see images now, or after one reload into the vision tower. */
+    get supportsImageInput(): boolean;
+    /**
+     * Cache adapter (Phase 1+). Set by ChatManager after construction.
+     * Phase 2 wires the real `MlxCacheAdapter` that talks to the
+     * wrapped `gezel_mlx_server.py`. Sessions call its
+     * `buildRequestExtras` on every send to derive `cache_id`.
+     */
+    private cacheAdapter;
+    setCacheAdapter(adapter: import('./cache-adapter.js').MlxCacheAdapter): void;
+    getCacheAdapter(): import('./cache-adapter.js').MlxCacheAdapter | null;
+    /**
+     * Current engine base URL — `null` when the supervised engine
+     * hasn't started yet. Cache adapter polls treat null as "not
+     * ready" and return empty stats.
+     */
+    currentBaseUrl(): string | null;
+    private lastTaggedPrefill;
+    /**
+     * The most recent prefill this engine ran for a DIFFERENT session, if it
+     * happened after `sinceMs`. Proof that the engine was alive and working —
+     * just not on the caller's turn.
+     */
+    prefillForOtherSessionSince(ownCacheId: string | null, sinceMs: number): {
+        detail: string;
+        at: number;
+    } | null;
+    /** True only while this sidecar is serving or queuing a physical request. */
+    isEngineBusy(): boolean;
+    /** Hold startup demand through the first request, before the engine gate exists. */
+    withEngineTurn<T>(run: () => Promise<T>, signal?: AbortSignal): Promise<T>;
+    engineLifecycleSnapshot(): NativeEngineLifecycleSnapshot | undefined;
+    acquireExclusiveEngineRequest(label: string, signal?: AbortSignal, onWait?: (info: {
+        aheadOf: number;
+    }) => void, capacityProtected?: boolean, priority?: EngineRequestPriority): Promise<() => void>;
+    /** Cache warms through the adapter: nobody is waiting on them, so they queue as background. */
+    runExclusiveEngineRequest<T>(label: string, work: () => Promise<T>): Promise<T>;
+    initialize(): Promise<void>;
+    getContextWindow(): number;
+    /** See `LLMProvider.plannedReservationBytes`. */
+    plannedReservationBytes(): number | undefined;
+    /**
+     * True once {@link shutdown} has been called. A session whose stream dies
+     * reads this to tell "Gezel stopped this engine" from "the engine died on
+     * its own" — the two need different messages, and the flag is set before
+     * the stop so the racing turn always sees it.
+     */
+    get isDisposed(): boolean;
+    shutdown(): Promise<void>;
+    createSession(opts: SessionOpts): Promise<LLMSession>;
+    /** Classifier hook wired into the supervisor's `onRawLine`. */
+    onStdoutLine(line: string): void;
+    /** See {@link EngineLogRouter.takeFatalError}. */
+    takeFatalError(): MlxFatalError | null;
+    /** See {@link EngineLogRouter.takeRuntimeError}. */
+    takeRuntimeError(): MlxFatalError | null;
+    _registerActiveSession(session: MlxSession): void;
+    _deregisterActiveSession(session: MlxSession): void;
+    listModels(): Promise<ModelInfo[]>;
+    getEffectiveModelId(): string | undefined;
+    private resolveBaseUrl;
+}
+interface MlxSessionDeps {
+    resolveBaseUrl: () => Promise<string>;
+    markUsed: () => void;
+    fetchImpl: typeof fetch;
+    model: string;
+    numCtx: number;
+    systemMessage: string;
+    /** Layered prefix-cache keys (flag ON only). See {@link SystemPromptLayers}. */
+    systemPromptLayers?: import('../../cache/adapter.js').SystemPromptLayers;
+    /** See {@link SessionOpts.systemSharedPrefix}. */
+    systemSharedPrefix?: string;
+    /** Volatile band seeded as a frozen system message after messages[0] (flag ON only). */
+    volatileContext?: string;
+    priorMessages: Array<{
+        role: 'user' | 'assistant';
+        content: string;
+        images?: string[];
+    } | {
+        role: 'assistant';
+        content: string;
+        toolCalls: ExternalToolCall[];
+    } | {
+        role: 'tool';
+        content: string;
+        toolCallId: string;
+    }>;
+    bridges: ToolExecutor;
+    /**
+     * Caller-supplied external tool definitions. Advertised to mlx-lm's
+     * OpenAI-compatible endpoint alongside any bridge tools; calls
+     * against these names halt the loop rather than executing — captured
+     * via {@link LLMSession.capturedToolCalls}.
+     */
+    externalTools?: ExternalToolSpec[];
+    queue: ProviderQueue;
+    provider: MlxProvider;
+    debug?: {
+        isEnabled(): boolean;
+    };
+    /**
+     * Mid-tool-loop compaction hook — see {@link SessionOpts.requestCompaction}.
+     * Same shape as the llama-cpp + ollama paths; the manager wires it
+     * to a one-shot synthesis call when the active provider is local.
+     */
+    requestCompaction?: NonNullable<SessionOpts['requestCompaction']>;
+    /**
+     * Resolved per-model behavior profile. Same role as Ollama's: gates
+     * the ramble detector + pre-tool preamble fold on profile opt-in.
+     */
+    profile?: ResolvedModelProfile;
+    /**
+     * Whether this model's chat template opens a reasoning block as the last
+     * thing it emits, so the stream starts mid-thought with only a closing tag
+     * to come. Read from the model directory at build time — guessing it from
+     * the token stream is exactly the ambiguity that made reasoning leak.
+     * Resolve per request with `requestOpensReasoning`: a switch-aware template
+     * opens nothing when the request turns thinking off.
+     */
+    templateOpensReasoning?: TemplateReasoningOpen;
+    /** Active craftbook step — passed to anti-spin abort messages. */
+    activeCraftbookStep?: NonNullable<SessionOpts['activeCraftbookStep']>;
+    /**
+     * Resolved sampling / reasoning / output / tool-choice knobs. Built
+     * by ChatManager from `gezel.tuning` + `catalog.manifest.tuning` and
+     * applied to the request body via `MLX_TUNING_MAP`. When unset, the
+     * provider sends only the request shape's required fields and lets
+     * mlx-vlm fall back to its own defaults.
+     */
+    tuning?: import('../../model-profile/index.js').ResolvedTuning;
+    /** Manager-authoritative direct file-work classification. */
+    forceDirectFileWork?: boolean;
+    /** See {@link SessionOpts.singleToolCallTurn}. */
+    singleToolCallTurn?: boolean;
+    terminalToolPolicy?: NonNullable<SessionOpts['terminalToolPolicy']>;
+    /** See {@link SessionOpts.onWireTranscript}. */
+    onWireTranscript?: NonNullable<SessionOpts['onWireTranscript']>;
+}
+declare class MlxSession extends StreamingSessionBase implements LLMSession {
+    private readonly deps;
+    get supportsImageInput(): boolean;
+    private readonly messages;
+    /**
+     * Direct provider callers and ephemeral one-shot work do not always have a
+     * persisted chat-session id to put in `opts.queue.sessionId`. The wrapped
+     * MLX server treats a missing `cache_id` as an instruction to clear the KV
+     * cache after every request, which makes every tool-loop continuation pay a
+     * full re-prefill. Keep a stable id for this MlxSession as the fallback;
+     * normal chats still override it with their durable Gezel session id.
+     */
+    private readonly fallbackCacheId;
+    private fallbackCacheUsed;
+    private sendSeq;
+    /** See LlamaCppSession.currentTurnStartIdx — same role here. */
+    private currentTurnStartIdx;
+    private turnPolicy;
+    /**
+     * Aggregated `<think>` / `<reasoning>` text captured this turn,
+     * across every iteration of the tool loop. Reset at the top of each
+     * `sendAndWaitInner`. ChatManager reads it via
+     * `getLastTurnReasoning()` after the turn resolves and stashes the
+     * trace on the assistant message so the chat bubble can render it
+     * behind a collapsed expander.
+     */
+    private lastTurnReasoning;
+    /**
+     * Tool-call bodies the salvage layer couldn't parse on the most
+     * recent turn after the retry budget exhausted. Captured at the
+     * budget-exhausted branch and read by ChatManager via
+     * `getLastTurnAttemptedToolCalls()` to stash on the assistant
+     * message for the debug bundle. Reset at the top of `sendAndWaitInner`.
+     */
+    private lastTurnAttemptedToolCalls;
+    /**
+     * Watchdog reset hook owned by the in-flight `sendAndWaitInner`. Set
+     * just before the fetch fires; cleared in the turn's cleanup. While
+     * non-null, any engine-phase event arriving via `publishEnginePhase`
+     * counts as "the model is making progress" and resets the idle timer.
+     * Without this, large-prompt prefill (e.g. 30K tokens at 230 tok/s ≈
+     * 2 minutes) trips the pre-first-byte watchdog even though the
+     * python child is happily chunking through prefill — the SSE stream
+     * is silent until the very first generated token.
+     */
+    private currentTurnIdleReset;
+    /**
+     * Handler owned by the in-flight request for the engine's wave-scheduling
+     * markers (`waiting` / `paused` / `admitted`). Set beside
+     * {@link currentTurnIdleReset} and cleared with it.
+     */
+    private currentTurnEngineQueue;
+    private readonly engineHold;
+    /**
+     * Last prefill phase event seen during the current turn — captured
+     * for the abort message so the user sees "aborted at prefill 47% /
+     * 14336 of 29930 tokens" instead of just "no first byte." Cleared
+     * at the start of every send.
+     */
+    private lastPrefillEvent;
+    /**
+     * The `cache_id` this session sends to the engine. Set on every send, and
+     * the key the provider routes per-sub engine phases by — see
+     * {@link MlxSession.ownsCacheId}.
+     */
+    private currentCacheId;
+    private readonly externalToolNames;
+    private capturedCalls;
+    private toolSurface;
+    /**
+     * Session-scoped quarantine for the assisted MTP route. Once an emitted
+     * Hermes call proves that decode-time required-field enforcement diverged,
+     * every later grammar-armed request from this live session uses ordinary
+     * sequential BatchGenerator decoding. The latch preserves speculation for
+     * healthy sessions while ensuring one bad call cannot become a five-retry
+     * validator loop.
+     */
+    private forceSequentialToolGrammar;
+    constructor(deps: MlxSessionDeps);
+    capturedToolCalls(): ExternalToolCall[];
+    getWireTranscript(): WireTranscriptEntry[] | undefined;
+    getToolSurface(): ToolSurfaceSize | undefined;
+    get numCtx(): number;
+    estimatePromptChars(): number;
+    /**
+     * Prefill the exact standing `[system][tools][transcript]` prefix without
+     * appending a user turn or mutating this session. Used by the machine broker
+     * after a user daemon prepares `/v1/remote/cache/warm`.
+     */
+    prefillOnly(opts?: {
+        timeoutMs?: number;
+        sessionId?: string;
+    }): Promise<void>;
+    getLastTurnReasoning(): string | undefined;
+    /**
+     * Running length of the same trace {@link getLastTurnReasoning}
+     * returns — read when a tool call fires so the persisted call can
+     * carry the offset it fired at.
+     */
+    getCurrentTurnReasoningLength(): number;
+    getLastTurnAttemptedToolCalls(): Array<{
+        body: string;
+        reason?: string;
+    }> | undefined;
+    /**
+     * Whether a phase tagged with this engine cache id belongs to this session.
+     * The id is the one sent as `cache_id` on every request, so it is exactly
+     * what the sidecar stamps onto its per-sub prefill markers.
+     */
+    ownsCacheId(cacheId: string): boolean;
+    publishEnginePhase(ev: EnginePhaseEvent): void;
+    publishEngineStats(ev: EngineStatsEvent): void;
+    private maybeCompactMidLoop;
+    get isDisposed(): boolean;
+    sendAndWait(prompt: string, opts?: SendAndWaitOpts): Promise<string>;
+    private sendAndWaitInner;
+    providerState(): ProviderSessionState;
+    getRegisteredToolNames(): string[];
+    setSystemMessage(text: string): void;
+    disconnect(): Promise<void>;
+}