mlx-bun 0.4.0 → 1.0.0

package/src/backends/mlx/speculative-round.ts

@@ -1,161 +0,0 @@
-import type { MlxArray } from "../../mlx/array";
-import * as ops from "../../mlx/ops";
-import type { Cancellation, GenerationOutput, TokenLogprobs } from "../../contracts/generation";
-import type { DisposableResource } from "../../contracts/resources";
-import { cleanupFailure, disposeResources } from "../../engine/resources";
-import type { SpeculativeTransaction } from "../../inference/rollback";
-import { DraftAcceptance } from "../../inference/draft-acceptance";
-import { deliverDraftOutputs, type DraftOutputResult } from "../../inference/draft-output";
-import { readStepExtras, type NumberStepSampler } from "../../sampler";
-
-/** Draft state owns proposal generation and its committed hidden/KV state.
- * It borrows verification context and knows nothing about target storage. */
-export interface MlxDraftRows<Context = MlxArray> {
-  /** Each row returns 0..depth proposals. The method verifies the returned
-   * lengths; an empty row still consumes its pending token. */
-  draft(pending: readonly number[], depth: number, steps: readonly number[]): number[][] | Promise<number[][]>;
-  commit(accepted: readonly number[], context: Context): void | Promise<void>;
-}
-
-/** A bound graph transfers both tensors on success and releases partial
- * allocations on failure. State retention is a separate injected interface. */
-export interface MlxRowVerification<Context extends DisposableResource = MlxArray> {
-  readonly transaction: SpeculativeTransaction<readonly number[]>;
-  forward(ids: MlxArray): Promise<{ logits: MlxArray; context: Context }>;
-}
-
-export interface MlxSpeculativeRow {
-  readonly pending: number;
-  readonly step: number;
-  readonly remaining: number;
-  readonly eosTokenIds: readonly number[];
-  readonly sampling: NumberStepSampler;
-  readonly grammarDone?: () => boolean;
-}
-
-export interface MlxSpeculativeRowResult {
-  readonly drafts: readonly number[];
-  readonly acceptance: DraftAcceptance;
-  /** One entry per content token; stopping EOS is excluded. */
-  readonly logprobs: readonly (TokenLogprobs | undefined)[];
-}
-
-export interface MlxSpeculativeOutputRow extends MlxSpeculativeRow {
-  readonly output: Pick<GenerationOutput, "commit">;
-  readonly cancellation?: Cancellation;
-}
-
-/** Execute one method step through output delivery and aligned state commit.
- * The caller retires terminal rows and publishes eligible completed state
- * before changing membership. Consumer failures are returned per request;
- * a graph/state failure throws and requires discarding the whole group. */
-export async function advanceSpeculativeOutputs<Context extends DisposableResource = MlxArray>(
-  rows: readonly MlxSpeculativeOutputRow[], depth: number,
-  draft: MlxDraftRows<NoInfer<Context>>, target: MlxRowVerification<Context>,
-): Promise<{ rounds: readonly MlxSpeculativeRowResult[]; outputs: readonly DraftOutputResult[] }> {
-  const round = await prepareSpeculativeRows(rows, depth, draft, target);
-  let outputs: DraftOutputResult[];
-  try {
-    outputs = await deliverDraftOutputs(round.results.map((result, row) => ({
-      acceptance: result.acceptance, logprobs: result.logprobs,
-      output: rows[row]!.output, cancellation: rows[row]!.cancellation,
-    })));
-    await round.commit(outputs.map(output => output.accepted));
-  } catch (error) { return cleanupFailure(error, () => round.dispose()); }
-  round.dispose();
-  return { rounds: round.results, outputs };
-}
-
-/** Retention stays unresolved while the method delivers output. A terminal
- * row may keep fewer accepted inputs for an output-aligned checkpoint;
- * continuing rows retain their complete accepted prefix. Commit exactly once,
- * then dispose; disposing an uncommitted round requires discarding its state. */
-export interface MlxPreparedSpeculativeRound extends DisposableResource {
-  readonly results: readonly MlxSpeculativeRowResult[];
-  commit(keepAccepted?: readonly number[]): Promise<void>;
-}
-
-/** One committed method step at B=1 or B>1. The caller supplies live rows with
- * a positive token budget and a common nonnegative draft depth. Admission,
- * cancellation boundaries, output delivery and cache publication stay outside
- * this step. On failure the caller discards the participating method state. */
-export async function advanceSpeculativeRows<Context extends DisposableResource = MlxArray>(
-  rows: readonly MlxSpeculativeRow[], depth: number,
-  draft: MlxDraftRows<NoInfer<Context>>, target: MlxRowVerification<Context>,
-): Promise<readonly MlxSpeculativeRowResult[]> {
-  const round = await prepareSpeculativeRows(rows, depth, draft, target);
-  try { await round.commit(); }
-  catch (error) { return cleanupFailure(error, () => round.dispose()); }
-  round.dispose();
-  return round.results;
-}
-
-/** Draft, verify and sample without deciding which output-aligned state to
- * retain. Graph-specific context is opaque to this operation. */
-export async function prepareSpeculativeRows<Context extends DisposableResource = MlxArray>(
-  rows: readonly MlxSpeculativeRow[], depth: number,
-  draft: MlxDraftRows<NoInfer<Context>>, target: MlxRowVerification<Context>,
-): Promise<MlxPreparedSpeculativeRound> {
-  const proposals = await draft.draft(rows.map(row => row.pending), depth, rows.map(row => row.step));
-  const B = rows.length;
-  const verifyDepth = proposals.reduce((max, tokens) => Math.max(max, tokens.length), 0);
-  using ids = ops.fromInt32(proposals.flatMap((tokens, row) => {
-    const window = [rows[row]!.pending, ...tokens];
-    // Causal right padding only fills the rectangular graph. Sampling stops
-    // at this row's own bonus position; commit discards every padded input.
-    while (window.length <= verifyDepth) window.push(rows[row]!.pending);
-    return window;
-  }), [B, verifyDepth + 1]);
-  target.transaction.begin(verifyDepth);
-  const verification = await target.forward(ids);
-  const dispose = () => disposeResources([verification.context, verification.logits]);
-  try {
-    const results = await sampleSpeculativeRows(rows, proposals, verification.logits);
-    return {
-      results,
-      async commit(accepted = results.map(result => result.acceptance.accepted)) {
-        target.transaction.resolve(accepted);
-        await draft.commit(accepted, verification.context);
-      },
-      dispose,
-    };
-  } catch (error) { return cleanupFailure(error, dispose); }
-}
-
-/** Sample a borrowed verify window. Each request owns its processor history,
- * grammar and RNG stream; completing one walk never samples its unused suffix.
- * The existing request adapter and grouped method use this same operation. */
-export async function sampleSpeculativeRows(
-  rows: readonly MlxSpeculativeRow[], proposals: readonly (readonly number[])[], logits: MlxArray,
-): Promise<MlxSpeculativeRowResult[]> {
-  const V = logits.shape[2]!, width = logits.shape[1]!;
-  const independent = rows[0]?.sampling.independent;
-  let tokens: number[] | undefined;
-  if (independent && rows.every(row => row.sampling.independent === independent)) {
-    using flat = ops.reshape(logits, [rows.length * width, V]);
-    using sampled = independent.sample(flat);
-    tokens = sampled.toIntTokens();
-  }
-  const results: MlxSpeculativeRowResult[] = [];
-  for (let row = 0; row < rows.length; row++) {
-    const request = rows[row]!;
-    const acceptance = new DraftAcceptance(proposals[row]!, request.remaining, request.eosTokenIds);
-    const logprobs: (TokenLogprobs | undefined)[] = [];
-    while (!acceptance.done) {
-      const position = acceptance.position;
-      let token: number, metadata: TokenLogprobs | undefined;
-      if (tokens) token = tokens[row * width + position]!;
-      else {
-        using slice = logits.slice([row, position, 0], [row + 1, position + 1, V]);
-        using scores = ops.reshape(slice, [1, V]);
-        const sampled = await request.sampling.sample(scores, request.step + acceptance.emitted.length);
-        token = sampled.token;
-        metadata = readStepExtras(sampled.extras);
-      }
-      acceptance.accept(token, request.grammarDone?.() ?? false);
-      if (!acceptance.sawEos) logprobs.push(metadata);
-    }
-    results.push({ drafts: proposals[row]!, acceptance, logprobs });
-  }
-  return results;
-}

package/src/backends/mlx/speculative.ts

@@ -1,114 +0,0 @@
-import type { MlxArray } from "../../mlx/array";
-import * as ops from "../../mlx/ops";
-import type { Cache } from "../../model/gemma4";
-import type { RuntimeModel } from "../../model/factory";
-import type { Sampler } from "../../sampler";
-import type { DraftProvider, DraftSource } from "../../spec/source";
-import { runtimeConfig, type RuntimeConfig } from "../../runtime-config";
-import type { MlxModelMemory } from "./autoregressive";
-import type { SpeculativeTransaction } from "../../inference/rollback";
-import { bindCacheRollback } from "./rollback";
-import type { GraphDescriptor } from "../../inference/graph";
-import { bindLegacyDraftTarget } from "./draft-target";
-
-/** Bound target operations and draft construction for one speculative run.
- * A replacement graph supplies this entire port, including any hidden taps.
- * The verifier never needs a RuntimeModel or an artifact-family check. */
-export interface MlxSpeculativeBinding {
-  readonly runtime?: RuntimeConfig;
-  readonly memory?: MlxModelMemory;
-  readonly descriptor: GraphDescriptor;
-  readonly eosTokenIds: readonly number[];
-  readonly prefillTailSplit: boolean;
-  makeCache(): Cache[];
-  openDraft(sampler: Sampler, caches: Cache[]): DraftSource;
-  bindRollback(caches: Cache[]): SpeculativeTransaction;
-  forward(ids: MlxArray, caches: Cache[], tapLayers?: number[]):
-    Promise<{ hidden: MlxArray; ctxML: MlxArray | null }>;
-  projectLogits(hidden: MlxArray): MlxArray;
-  /** Establish the implementation's verify kernel context; restore on close. */
-  pinVerify?(): { close(): void };
-}
-
-/** Legacy mutable tap/kernel fields require the gateway's exclusive lease. */
-export function bindLegacySpeculativeModel(model: RuntimeModel, provider: DraftProvider): MlxSpeculativeBinding {
-  const runtime = runtimeConfig();
-  return {
-    runtime,
-    memory: model,
-    descriptor: Object.freeze({ id: `legacy-spec:${model.config.modelType}`, backend: "mlx",
-      graphAbi: "mlx-hidden-bsh-v1", stateAbi: "legacy-cache-array-v1", artifact: "legacy-resident-model" }),
-    eosTokenIds: model.config.eosTokenIds,
-    prefillTailSplit: runtime.flag("MLX_BUN_PREFILL_TAIL_SPLIT", true),
-    makeCache: model.makeCache.bind(model),
-    openDraft: (sampler, caches) => provider.open({ sampler, target: bindLegacyDraftTarget(model, caches) }),
-    bindRollback: bindCacheRollback,
-    forward: (ids, caches, tapLayers) => legacyForwardWithTaps(model, ids, caches, tapLayers),
-    projectLogits: model.logitsFromHidden.bind(model),
-    ...("setSpecKernelPinned" in model ? {
-      pinVerify() {
-        model.setSpecKernelPinned(true);
-        return { close() { model.setSpecKernelPinned(false); } };
-      },
-    } : {}),
-  };
-}
-
-export function assertMlxSpeculativeBinding(binding: MlxSpeculativeBinding): void {
-  const descriptor = binding.descriptor;
-  if (descriptor.backend !== "mlx" || descriptor.graphAbi !== "mlx-hidden-bsh-v1" ||
-      descriptor.stateAbi !== "legacy-cache-array-v1")
-    throw new Error(`speculative binding ${descriptor.id} has an incompatible backend, graph, or state ABI`);
-}
-
-/** Run one target forward, and when `tapLayers` is set (⟹ a Gemma4 target — a
- *  DSpark-style source), also return the captured multi-layer context
- *  [1,L,m*H] (tapLayers concatenated on the feature axis; index nLayers is the
- *  post-finalNorm sentinel). Non-tapping sources get ctxML=null and never
- *  touch model.hiddenTap. Mirrors generate-dflash.ts forwardTapped. */
-async function legacyForwardWithTaps(
-  model: RuntimeModel,
-  ids: MlxArray,
-  caches: Cache[],
-  tapLayers: number[] | undefined,
-): Promise<{ hidden: MlxArray; ctxML: MlxArray | null }> {
-  if (!tapLayers) {
-    const asyncModel = model as RuntimeModel & {
-      forwardHiddenAsync?: (
-        ids: MlxArray,
-        caches: Cache[],
-      ) => Promise<MlxArray>;
-    };
-    const hidden = typeof asyncModel.forwardHiddenAsync === "function"
-      ? await asyncModel.forwardHiddenAsync(ids, caches)
-      : model.forwardHidden(ids, caches);
-    return { hidden, ctxML: null };
-  }
-  if (!("hiddenTap" in model)) throw new Error("target graph does not support hidden taps");
-  const m = model;
-  const previousTap = m.hiddenTap;
-  const cap = new Map<number, MlxArray>();
-  m.hiddenTap = { layers: new Set(tapLayers), captured: cap };
-  let hidden: MlxArray | null = null;
-  try {
-    hidden = model.forwardHidden(ids, caches);
-    const perLayer = tapLayers.map((li) => {
-      const a = cap.get(li);
-      if (!a) throw new Error(`spec tap: layer ${li} not captured`);
-      return a;
-    });
-    const ctxML = ops.concatAxis(perLayer, 2); // [1,L,m*H]
-    for (const [, a] of cap) a.dispose();
-    cap.clear(); // consumed — the finally must not double-dispose
-    const out = { hidden, ctxML };
-    hidden = null; // ownership returned to the caller
-    return out;
-  } finally {
-    // On any throw (forward mid-capture, missing layer, concat), free the
-    // partially-captured tap tensors and the orphaned hidden. On success both
-    // are already gone (cap cleared, hidden nulled).
-    hidden?.dispose();
-    for (const [, a] of cap) a.dispose();
-    m.hiddenTap = previousTap;
-  }
-}

package/src/backends/mlx/state-rows.ts

@@ -1,104 +0,0 @@
-import type { BatchableCache, Cache } from "../../model/gemma4-base";
-import { applyStateChanges, cleanupFailure, disposeResources } from "../../engine/resources";
-import type { PreparedStateChange } from "../../contracts/resources";
-
-/** Preserve the graph's layout capabilities when an empty sibling is made. */
-export interface MlxRowLayout<Layout extends BatchableCache> extends BatchableCache {
-  makeEmptyBatch(): Layout;
-}
-
-/** Layer layouts own tensor geometry. This collection only applies the same
- * row membership operation to every layer. It has no scheduling, sampling,
- * checkpoint-key or persistence policy. */
-export class MlxStateRows<Layout extends MlxRowLayout<Layout> = BatchableCache> {
-  /** Borrowed by the bound target graph; collection ownership stays here. */
-  readonly caches: Layout[];
-  #rows = 0;
-
-  /** Takes ownership of empty, compatible layer layouts. */
-  constructor(layouts: readonly Layout[]) { this.caches = [...layouts]; }
-  get rowCount(): number { return this.#rows; }
-  /** Retain an immutable batch boundary without extracting or materializing
-   * individual rows. Layer layouts decide how to share their backing arrays. */
-  clone(): MlxStateRows<Layout> {
-    const layouts: Layout[] = [];
-    try {
-      for (const source of this.caches) {
-        const layout = source.makeEmptyBatch(); layouts.push(layout);
-        layout.mergeRows([source]);
-      }
-      const result = new MlxStateRows(layouts); result.#rows = this.#rows;
-      return result;
-    } catch (error) { return cleanupFailure(error, () => disposeResources(layouts)); }
-  }
-  /** Prepare an initial cohort with one merge per layer. */
-  mergeRows(rows: readonly (readonly Cache[])[]): void {
-    const next: Layout[] = [];
-    try {
-      for (let layer = 0; layer < this.caches.length; layer++) {
-        const layout = this.caches[layer]!.makeEmptyBatch(); next.push(layout);
-        layout.mergeRows(rows.map(row => row[layer]!));
-      }
-    } catch (error) { return cleanupFailure(error, () => disposeResources(next)); }
-    const previous = this.caches.splice(0, this.caches.length, ...next);
-    this.#rows = rows.length; disposeResources(previous);
-  }
-
-  /** Borrow a prepared request's state. Build every replacement before
-   * releasing the active group, so a failed admission leaves it intact. */
-  append(row: readonly Cache[]): void {
-    applyStateChanges([() => this.prepareAppend(row)]);
-  }
-
-  /** Stage membership so a method can compose it with companion-state admission. */
-  prepareAppend(row: readonly Cache[]): PreparedStateChange {
-    return this.prepareAppendMany([row]);
-  }
-
-  /** Stage several admissions with one membership operation per layer. */
-  prepareAppendMany(rows: readonly (readonly Cache[])[]): PreparedStateChange {
-    const next: Layout[] = [];
-    try {
-      for (let layer = 0; layer < this.caches.length; layer++) {
-        const layout = this.caches[layer]!;
-        const merged = layout.makeEmptyBatch();
-        next.push(merged);
-        merged.mergeRows([...(this.#rows ? [layout] : []), ...rows.map(row => row[layer]!)]);
-      }
-    } catch (error) { return cleanupFailure(error, () => disposeResources(next)); }
-    return {
-      commit: () => {
-        const previous = this.caches.splice(0, this.caches.length, ...next.splice(0));
-        this.#rows += rows.length;
-        next.push(...previous);
-      },
-      dispose: () => disposeResources(next.splice(0)),
-    };
-  }
-
-  /** The method calls this only between committed steps. A layer failure
-   * invalidates the group; its caller discards the state. */
-  filterRows(keep: readonly number[]): void {
-    if (!keep.length) { this.clear(); return; }
-    for (const layout of this.caches) layout.filterRows(keep);
-    this.#rows = keep.length;
-  }
-
-  /** Returns independent request state for a cache publisher or a new group. */
-  extractRow(row: number): Cache[] {
-    const state: Cache[] = [];
-    try { for (const layout of this.caches) state.push(layout.extractRow(row)); return state; }
-    catch (error) { return cleanupFailure(error, () => disposeResources(state)); }
-  }
-
-  clear(): void {
-    const next: Layout[] = [];
-    try { for (const layout of this.caches) next.push(layout.makeEmptyBatch()); }
-    catch (error) { return cleanupFailure(error, () => disposeResources(next)); }
-    const previous = this.caches.splice(0, this.caches.length, ...next);
-    this.#rows = 0;
-    disposeResources(previous);
-  }
-
-  dispose(): void { this.#rows = 0; disposeResources(this.caches); }
-}

package/src/backends/mlx/state-views.ts

@@ -1,46 +0,0 @@
-import type { StateView } from "../../contracts/resources";
-import { disposeResources, ownResource } from "../../engine/resources";
-import type { MlxArray } from "../../mlx/array";
-import { isBatchableCache, type Cache } from "../../model/gemma4-base";
-import { cloneKvCaches, type CacheCodecProvider } from "../../kv-store";
-
-/** Compatibility boundary for legacy state(). Numerical code consumes a lease,
- * never the legacy ownership marker. No copy, evaluation or fence is added. */
-export function leaseCacheState(cache: Cache): StateView<MlxArray> {
-  return ownResource<readonly MlxArray[]>(cache.state(), cache.stateNeedsDispose ? disposeResources : () => {});
-}
-
-/** Preserve the existing prefill evaluation order: temporary views first,
- * borrowed state second. Partial acquisition releases every acquired view. */
-export function leaseCacheStates(caches: readonly Cache[]): StateView<MlxArray> {
-  const leases: StateView<MlxArray>[] = [];
-  const close = () => disposeResources(leases.map((lease) => ({ dispose: () => lease.close() })));
-  try {
-    for (const owned of [true, false]) {
-      for (const cache of caches) if (!!cache.stateNeedsDispose === owned) leases.push(leaseCacheState(cache));
-    }
-    return ownResource(leases.flatMap((lease) => [...lease.borrow()]), close);
-  } catch (error) {
-    try { close(); } catch (cleanupError) { throw new AggregateError([error, cleanupError], "state view acquisition failed"); }
-    throw error;
-  }
-}
-
-/** Precision/storage transitions constrain reuse independently of scheduling. */
-export function minimumReusableOffset(caches: readonly Pick<Cache, "minimumReusableOffset">[]): number {
-  return Math.max(0, ...caches.map(cache => cache.minimumReusableOffset ?? 0));
-}
-
-/** Snapshot an adopted single-row state into persistence-compatible caches.
- * Row layouts own compact extraction; ordinary serial caches use their codec. */
-export function cloneSingleRowState(caches: readonly Cache[], codecs?: CacheCodecProvider): Cache[] {
-  const owned: Cache[] = [];
-  try {
-    for (const cache of caches) owned.push(...(isBatchableCache(cache)
-      ? [cache.extractRow(0)] : cloneKvCaches([cache], codecs)));
-    return owned;
-  } catch (error) {
-    try { disposeResources(owned); } catch (cleanupError) { throw new AggregateError([error, cleanupError], "single-row snapshot failed"); }
-    throw error;
-  }
-}

package/src/backends/mlx/target-layout-capability.ts

@@ -1,35 +0,0 @@
-import { Glm52Cache } from "../../model/glm52-cache";
-import { SpeculativeRotatingKVCache } from "../../model/speculative-rotating-kv";
-import { RotatingAffineLayout, SpeculativeRotatingAffineLayout } from "../../model/rotating-kv-layout";
-import type { Cache } from "../../model/gemma4-base";
-import { KVCache, QuantizedKVCache, TurboQuantKVCache, RotatingKVCache, RotatingQuantizedKVCache } from "../../model/gemma4-base";
-import { DelayedRotatingQuantizedKVCache } from "../../model/delayed-rotating-quantized-kv";
-import { DelayedQuantizedKVCache } from "../../model/delayed-quantized-kv";
-import { DelayedTurboQuantKVCache } from "../../model/delayed-turboquant-kv";
-import { SSMCache } from "../../model/qwen3-delta";
-import { BatchedSSMCache } from "../../model/batched-ssm";
-import { BatchedKVCache } from "../../model/batched-kv";
-import { BatchedQuantizedKVCache } from "../../model/batched-quantized-kv";
-import { BatchedTurboQuantKVCache } from "../../model/batched-turboquant-kv";
-
-export type TargetRowLayout = Glm52Cache | DelayedRotatingQuantizedKVCache | RotatingAffineLayout | SpeculativeRotatingKVCache | BatchedKVCache | BatchedQuantizedKVCache | BatchedTurboQuantKVCache |
-  BatchedSSMCache | DelayedTurboQuantKVCache | DelayedQuantizedKVCache;
-
-/** The storage binding advertises a row transaction factory. The gateway can
- * compose supported model graphs without identifying their model family. */
-export function targetRowLayoutFactory(cache: Cache): (() => TargetRowLayout) | undefined {
-  if (cache instanceof Glm52Cache) return () => cache.makeEmptyBatch();
-  if (cache instanceof DelayedRotatingQuantizedKVCache)
-    return () => new DelayedRotatingQuantizedKVCache(cache.maxSize, cache.groupSize, cache.bits, cache.start, cache.maintain, undefined, true);
-  if (cache instanceof RotatingAffineLayout || cache instanceof RotatingQuantizedKVCache)
-    return () => new SpeculativeRotatingAffineLayout(cache.maxSize, cache.groupSize, cache.bits);
-  if (cache instanceof RotatingKVCache || cache instanceof SpeculativeRotatingKVCache) return () => new SpeculativeRotatingKVCache(cache.maxSize);
-  if (cache instanceof DelayedTurboQuantKVCache || cache instanceof DelayedQuantizedKVCache ||
-      cache instanceof BatchedTurboQuantKVCache || cache instanceof BatchedKVCache ||
-      cache instanceof BatchedQuantizedKVCache || cache instanceof BatchedSSMCache) return () => cache.makeEmptyBatch();
-  if (cache instanceof TurboQuantKVCache) return () => new BatchedTurboQuantKVCache(cache.kBits, cache.vBits);
-  if (cache instanceof SSMCache) return () => new BatchedSSMCache();
-  if (cache instanceof QuantizedKVCache) return () => new BatchedQuantizedKVCache(cache.groupSize, cache.bits);
-  if (cache instanceof KVCache) return () => new BatchedKVCache();
-  return undefined;
-}

package/src/backends/mlx/text-engine.ts

@@ -1,73 +0,0 @@
-import { createCompletionMethod } from "../../serve/session-completion-engine";
-import type { ExecutionPlanner, InferenceEngine, Timer } from "../../contracts/generation";
-import { createInferenceEngine, type EngineOptions } from "../../engine/engine";
-import { throwIfCancelled } from "../../engine/cancellation";
-import type { GenerateOptions, GenerateStats } from "../../generate";
-import type { GenerationGateway, RequestShape } from "../../serve/generation-gateway";
-
-/** Data-only API for text requests whose policy/shape was already
- * resolved by the caller. Native request resources still belong to the prepared
- * completion API; they are deliberately absent from this token API. */
-export type TextGenerationOptions = Omit<GenerateOptions,
-  "signal" | "cache" | "grammar" | "fill" | "adapters" | "promptEmbeddings" |
-  "imageMask" | "multimodalMask" | "visionPixels" | "onPrefillDone" |
-  "onDecodeCheckpoint" | "initialPendingToken" | "initialGeneratedTokens" |
-  "originalPromptTokens" | "snapshotAt" | "checkpointEveryTokens"
->;
-
-export interface TextGenerationRequest {
-  readonly promptIds: readonly number[];
-  readonly options: TextGenerationOptions;
-  readonly shape: RequestShape;
-}
-
-export const hostTimer: Timer = {
-  after(milliseconds, callback) {
-    const timer = setTimeout(callback, milliseconds);
-    return () => { clearTimeout(timer); };
-  },
-};
-
-/** Reuses the gateway's placement, exclusive lease, prompt cache, batching,
- * and quant-specific execution. The portable session never imports them.
- * Closing this engine cancels its sessions; it does not dispose a shared gateway. */
-export function createTextInferenceEngine(
-  gateway: Pick<GenerationGateway, "place" | "run">,
-  options: Omit<EngineOptions, "timer"> & { timer?: Timer } = {},
-): InferenceEngine<TextGenerationRequest, GenerateStats> {
-  const planner: ExecutionPlanner<TextGenerationRequest, GenerateStats> = {
-    async plan(request, cancellation) {
-      throwIfCancelled(cancellation);
-      if (request.shape.hasVision || request.shape.hasAdapters || request.shape.hasGrammar)
-        throw new Error("native request resources require the prepared completion API");
-      const placement = gateway.place(request.shape, request.options);
-      const maxTokens = request.options.maxTokens ?? 512;
-      return {
-        id: `text:${placement.mechanism}`, outputTokenLimit: maxTokens,
-        method: createCompletionMethod(gateway, {
-          prompt: [...request.promptIds], options: request.options,
-          shape: request.shape, placement,
-        }),
-      };
-    },
-  };
-  const engine = createInferenceEngine(planner, { ...options, timer: options.timer ?? hostTimer });
-  return {
-    open(request, control) {
-      // Snapshot before demand-start: callers may reuse/mutate request arrays.
-      // TextGenerationOptions contains data only; no native handles or callbacks.
-      const snapshot = structuredClone(request);
-      if (!snapshot.promptIds.length || snapshot.promptIds.some((id) => !Number.isInteger(id) || id < 0 || id > 0xffffffff))
-        return Promise.reject(new Error("prompt must contain valid token ids"));
-      const maxTokens = snapshot.options.maxTokens ?? 512;
-      if (!Number.isSafeInteger(maxTokens) || maxTokens <= 0)
-        return Promise.reject(new Error("text generation requires a positive token limit"));
-      if ((snapshot.options.topLogprobs ?? 0) > (options.maxTopLogprobs ?? 20))
-        return Promise.reject(new Error("top logprobs exceed session delivery capacity"));
-      if (control.output === "collect" && (snapshot.options.logprobs || snapshot.options.topLogprobs))
-        return Promise.reject(new Error("token logprobs require stream consumption"));
-      return engine.open(snapshot, control);
-    },
-    close: () => engine.close(),
-  };
-}