mlx-bun 0.0.13 → 0.2.0

package/src/serve/completion-executor.ts

@@ -0,0 +1,346 @@
+import type {
+  GenerateOptions,
+  GenerateStats,
+  TokenLogprobs,
+} from "../generate";
+import type {
+  GenerationMechanism,
+  GenerationPlacement,
+  OnToken,
+  RequestShape,
+  Vision,
+} from "./generation-gateway";
+import {
+  type CompletionEvent,
+  type CompletionFinishReason,
+  type CompletionFlowControl,
+  type CompletionResult,
+  type CompletionSinkOptions,
+  CompletionSink,
+  type CompletionToolCall,
+} from "./completion-sink";
+import { recordLane, type Lane } from "./lane-registry";
+import {
+  planRequest,
+  type RequestPlan,
+  type RequestPlanInput,
+  type RequestRejection,
+} from "./request-plan";
+import type { PromptResponseTrace } from "./prompt-response-trace";
+
+export interface CompletionEngine {
+  place(shape: RequestShape): GenerationPlacement;
+  run(
+    promptIds: number[],
+    options: GenerateOptions & { stopSequences?: string[] },
+    onToken: OnToken,
+    vision: Vision | undefined,
+    shape: RequestShape,
+    placement: GenerationPlacement,
+    signal?: AbortSignal,
+    trace?: PromptResponseTrace,
+  ): Promise<GenerateStats>;
+}
+
+export interface CompletionPreparation {
+  requestId: string;
+  plan: RequestPlanInput;
+  vision?: Vision;
+  pipeline: Omit<CompletionSinkOptions, "flowControl">;
+  createFlowControl?(placement: CompletionPlacement):
+    | CompletionSinkOptions["flowControl"];
+  onPlacement?(placement: CompletionPlacement): void;
+  idToToken(id: number): string;
+}
+
+export interface CompletionPlacement {
+  mechanism: GenerationMechanism;
+  lane: Lane;
+  shape: RequestShape;
+}
+
+interface PreparedData extends Omit<CompletionPreparation, "plan"> {
+  plan: RequestPlan;
+  stream: boolean;
+  requestedMaxTokens: number;
+}
+
+const preparedBrand: unique symbol = Symbol("PreparedCompletion");
+
+/** A single-use completion declaration. Protocol adapters can create one but
+ * cannot inspect or rewrite the resolved execution plan after preparation. */
+export interface PreparedCompletion {
+  readonly [preparedBrand]: true;
+}
+
+const preparedData = new WeakMap<PreparedCompletion, PreparedData>();
+
+export function prepareCompletion(data: CompletionPreparation): PreparedCompletion {
+  const planned = planRequest(data.plan);
+  if (!planned.ok) {
+    planned.dispose();
+    throw new CompletionRejected(planned);
+  }
+  const prepared = Object.freeze({ [preparedBrand]: true }) as PreparedCompletion;
+  preparedData.set(prepared, {
+    requestId: data.requestId,
+    plan: planned,
+    stream: data.plan.stream,
+    requestedMaxTokens: data.plan.requestedMaxTokens,
+    ...(data.vision ? { vision: data.vision } : {}),
+    pipeline: data.pipeline,
+    ...(data.createFlowControl
+      ? { createFlowControl: data.createFlowControl }
+      : {}),
+    ...(data.onPlacement ? { onPlacement: data.onPlacement } : {}),
+    idToToken: data.idToToken,
+  });
+  return prepared;
+}
+
+export interface CompletionControl {
+  signal?: AbortSignal;
+  trace?: PromptResponseTrace;
+  onEvents?(events: readonly CompletionEvent[]):
+    | void
+    | false
+    | Promise<void | false>;
+  /** Receives one live, read-only usage view before execution. Retain the
+   *  reference to close an errored protocol stream with accumulated usage.
+   *  The executor updates it in place, avoiding per-token callback/allocation
+   *  work on the decode path. */
+  onUsageProgress?(usage: Readonly<CompletionUsage>): void;
+}
+
+export interface CompletionUsage {
+  promptTokens: number;
+  cachedTokens: number;
+  completionTokens: number;
+  totalTokens: number;
+  speculation?: GenerateStats["spec"];
+}
+
+export interface CompletionSummary {
+  content: string;
+  reasoning: string;
+  toolCalls: CompletionToolCall[];
+  stopped: boolean;
+  finishReason: CompletionFinishReason;
+  lane: Lane;
+  usage: CompletionUsage;
+  logprobs: { content: Record<string, unknown>[] } | null;
+}
+
+export class CompletionRejected extends Error {
+  readonly status: number;
+  readonly error: RequestRejection["error"];
+
+  constructor(rejection: RequestRejection) {
+    super(rejection.error.message);
+    this.name = "CompletionRejected";
+    this.status = rejection.status;
+    this.error = rejection.error;
+  }
+}
+
+class LogprobsCollector {
+  readonly #tokens: number[] = [];
+  readonly #tokenLogprobs: number[] = [];
+  readonly #topTokens: { id: number; token: string; logprob: number }[][] = [];
+
+  constructor(
+    private readonly wantLogprobs: boolean,
+    private readonly topK: number,
+    private readonly idToToken: (id: number) => string,
+  ) {}
+
+  push(token: number, info?: TokenLogprobs): void {
+    this.#tokens.push(token);
+    if (this.wantLogprobs) this.#tokenLogprobs.push(info?.logprob ?? NaN);
+    if (this.topK > 0) {
+      this.#topTokens.push(
+        (info?.top ?? []).map((candidate) => ({
+          id: candidate.id,
+          token: this.idToToken(candidate.id),
+          logprob: candidate.logprob,
+        })),
+      );
+    }
+  }
+
+  payload(): { content: Record<string, unknown>[] } | null {
+    if (this.#topTokens.length) {
+      return {
+        content: this.#topTokens.map((tokens) =>
+          tokens.length ? { ...tokens[0]!, top_logprobs: tokens } : {},
+        ),
+      };
+    }
+    if (this.#tokenLogprobs.length) {
+      return {
+        content: this.#tokens.map((id, index) => ({
+          id,
+          logprob: this.#tokenLogprobs[index]!,
+        })),
+      };
+    }
+    return null;
+  }
+}
+
+function initialLane(mechanism: GenerationMechanism, shape: RequestShape): Lane {
+  if (mechanism === "continuous") return "batched";
+  return shape.hasDraft ? "serial+spec" : "serial";
+}
+
+function finalLane(
+  mechanism: GenerationMechanism,
+  lane: Lane,
+  stats: GenerateStats,
+): Lane {
+  if (mechanism === "continuous") return "batched";
+  return stats.spec ? "serial+spec" : lane;
+}
+
+function finishReason(
+  result: CompletionResult,
+  stats: GenerateStats,
+  maxTokens: number,
+): CompletionFinishReason {
+  if (result.toolCalls.length) return "tool_calls";
+  if (result.stopped) return "stop";
+  return stats.generatedTokens >= maxTokens ? "length" : "stop";
+}
+
+function combineControl(
+  sinkControl: CompletionFlowControl,
+  eventControl: void | false | Promise<void | false>,
+): void | false | Promise<void | false> {
+  if (sinkControl === false || eventControl === false) return false;
+  if (!(sinkControl instanceof Promise) && !(eventControl instanceof Promise))
+    return undefined;
+  return Promise.all([sinkControl, eventControl]).then((values) =>
+    values.includes(false) ? false : undefined,
+  );
+}
+
+/** Owns one completion attempt from admission through terminal accounting.
+ * Model math and scheduling policy remain in the supplied CompletionEngine. */
+export class CompletionExecutor {
+  constructor(private readonly engine: CompletionEngine) {}
+
+  async execute(
+    prepared: PreparedCompletion,
+    control: CompletionControl = {},
+  ): Promise<CompletionSummary> {
+    const input = preparedData.get(prepared);
+    if (!input) throw new Error("prepared completion has already been executed");
+    preparedData.delete(prepared);
+
+    const planned = input.plan;
+    const closeCompletion = control.trace?.begin("completion.total");
+    const usageProgress: CompletionUsage | null = control.onUsageProgress
+      ? {
+          promptTokens: planned.promptIds.length,
+          cachedTokens: 0,
+          completionTokens: 0,
+          totalTokens: planned.promptIds.length,
+        }
+      : null;
+
+    try {
+      if (usageProgress) control.onUsageProgress!(usageProgress);
+      const closePlacement = control.trace?.begin("completion.placement");
+      const enginePlacement = this.engine.place(planned.shape);
+      closePlacement?.();
+      if (enginePlacement.shape !== planned.shape)
+        throw new Error("generation placement does not belong to this request shape");
+      const mechanism = enginePlacement.mechanism;
+      const lane = initialLane(mechanism, planned.shape);
+      recordLane(input.requestId, lane);
+      const placement = { mechanism, lane, shape: planned.shape };
+      input.onPlacement?.(placement);
+      const flowControl = input.createFlowControl?.(placement);
+      const sink = new CompletionSink({
+        ...input.pipeline,
+        ...(flowControl ? { flowControl } : {}),
+      });
+      const logprobs = planned.captureLogprobs
+        ? new LogprobsCollector(
+            planned.wantLogprobs,
+            planned.topLogprobs,
+            input.idToToken,
+          )
+        : null;
+      control.signal?.throwIfAborted();
+      planned.transferOwnership();
+
+      const consumeToken: OnToken = (token, tokenLogprobs) => {
+        logprobs?.push(token, tokenLogprobs);
+        const pushed = sink.push(token);
+        const observed = pushed.events.length
+          ? control.onEvents?.(pushed.events)
+          : undefined;
+        return combineControl(pushed.control, observed);
+      };
+      const onToken: OnToken = usageProgress
+        ? (token, tokenLogprobs) => {
+            usageProgress.completionTokens++;
+            usageProgress.totalTokens++;
+            return consumeToken(token, tokenLogprobs);
+          }
+        : consumeToken;
+      const stats = await this.engine.run(
+        planned.promptIds,
+        planned.options,
+        onToken,
+        input.vision,
+        planned.shape,
+        enginePlacement,
+        control.signal,
+        control.trace,
+      );
+
+      if (input.stream) control.signal?.throwIfAborted();
+      const result = sink.finish();
+      if (result.events.length) await control.onEvents?.(result.events);
+
+      const resolvedLane = finalLane(mechanism, lane, stats);
+      if (resolvedLane !== lane) recordLane(input.requestId, resolvedLane);
+      if (usageProgress) {
+        usageProgress.promptTokens = stats.promptTokens;
+        usageProgress.cachedTokens = stats.cachedTokens;
+        usageProgress.completionTokens = stats.generatedTokens;
+        usageProgress.totalTokens = stats.promptTokens + stats.generatedTokens;
+        if (stats.spec) usageProgress.speculation = stats.spec;
+        else delete usageProgress.speculation;
+      }
+      const summary = {
+        content: result.content,
+        reasoning: result.reasoning,
+        toolCalls: result.toolCalls,
+        stopped: result.stopped,
+        finishReason: finishReason(
+          result,
+          stats,
+          planned.options.maxTokens ?? input.requestedMaxTokens,
+        ),
+        lane: resolvedLane,
+        usage: {
+          promptTokens: stats.promptTokens,
+          cachedTokens: stats.cachedTokens,
+          completionTokens: stats.generatedTokens,
+          totalTokens: stats.promptTokens + stats.generatedTokens,
+          ...(stats.spec ? { speculation: stats.spec } : {}),
+        },
+        logprobs: logprobs?.payload() ?? null,
+      };
+      closeCompletion?.();
+      return summary;
+    } catch (error) {
+      closeCompletion?.();
+      planned.dispose();
+      throw error;
+    }
+  }
+}

package/src/serve/completion-sink.ts

@@ -0,0 +1,157 @@
+export interface CompletionToolCall {
+  id: string;
+  type: "function";
+  function: { name: string; arguments: string };
+}
+
+export type CompletionEvent =
+  | { type: "content"; text: string }
+  | { type: "reasoning"; text: string }
+  | { type: "tool_calls"; calls: CompletionToolCall[] };
+
+export type CompletionFlowControl = false | Promise<void> | undefined;
+
+export interface CompletionPush {
+  events: CompletionEvent[];
+  control: CompletionFlowControl;
+}
+
+export interface CompletionResult {
+  events: CompletionEvent[];
+  content: string;
+  reasoning: string;
+  toolCalls: CompletionToolCall[];
+  stopped: boolean;
+}
+
+export type CompletionFinishReason = "stop" | "length" | "tool_calls";
+
+export interface CompletionStreamProtocol {
+  start(): string[];
+  addEvents(events: CompletionEvent[]): string[];
+  finish(reason: CompletionFinishReason, usage: Record<string, unknown>): string[];
+  error(message: string): string[];
+}
+
+export interface TokenTextRouter {
+  push(token: number): string;
+  flush(): string;
+  takeReasoning(): string;
+  toolCalls(): CompletionToolCall[];
+}
+
+export interface TextStopper {
+  readonly stopped: boolean;
+  push(text: string): string;
+  flush(): string;
+}
+
+export interface ThinkingSplitter {
+  push(text: string): { content: string; reasoning: string };
+  flush(): { content: string; reasoning: string };
+}
+
+export interface CompletionSinkOptions {
+  router: TokenTextRouter;
+  stopper: TextStopper;
+  thinking: ThinkingSplitter;
+  collectToolCalls: boolean;
+  flowControl?: (emitted: boolean) => CompletionFlowControl;
+}
+
+/** One generated-token sink for every text-serving surface.
+ *
+ * Chat uses the full router → stop matcher → thinking splitter stack.
+ * Text completions use the same sink with a plain router, thinking disabled,
+ * and tool-call collection off. `push()` returns both semantic events and the
+ * generation callback's control signal: false halts decode; a Promise yields
+ * before the next token. */
+export class CompletionSink {
+  readonly #router: TokenTextRouter;
+  readonly #stopper: TextStopper;
+  readonly #thinking: ThinkingSplitter;
+  readonly #collectToolCalls: boolean;
+  readonly #flowControl?: (emitted: boolean) => CompletionFlowControl;
+  #content = "";
+  #reasoning = "";
+
+  constructor(options: CompletionSinkOptions) {
+    this.#router = options.router;
+    this.#stopper = options.stopper;
+    this.#thinking = options.thinking;
+    this.#collectToolCalls = options.collectToolCalls;
+    this.#flowControl = options.flowControl;
+  }
+
+  get stopped(): boolean {
+    return this.#stopper.stopped;
+  }
+
+  push(token: number): CompletionPush {
+    const events: CompletionEvent[] = [];
+    const rawContent = this.#router.push(token);
+    this.#emitReasoning(events, this.#router.takeReasoning());
+    this.#emitParts(events, this.#thinking.push(this.#stopper.push(rawContent)));
+    return {
+      events,
+      control: this.#stopper.stopped
+        ? false
+        : this.#flowControl?.(events.length > 0),
+    };
+  }
+
+  finish(): CompletionResult {
+    const events: CompletionEvent[] = [];
+    if (!this.#stopper.stopped) {
+      const flushed = this.#router.flush();
+      this.#emitReasoning(events, this.#router.takeReasoning());
+      let tail = this.#stopper.push(flushed);
+      if (!this.#stopper.stopped) tail += this.#stopper.flush();
+      this.#emitParts(events, this.#thinking.push(tail));
+    }
+    this.#emitParts(events, this.#thinking.flush());
+
+    const toolCalls = this.#collectToolCalls ? this.#router.toolCalls() : [];
+    if (toolCalls.length) events.push({ type: "tool_calls", calls: toolCalls });
+    return {
+      events,
+      content: this.#content,
+      reasoning: this.#reasoning,
+      toolCalls,
+      stopped: this.#stopper.stopped,
+    };
+  }
+
+  #emitParts(
+    events: CompletionEvent[],
+    parts: { content: string; reasoning: string },
+  ): void {
+    this.#emitReasoning(events, parts.reasoning);
+    if (parts.content) {
+      this.#content += parts.content;
+      events.push({ type: "content", text: parts.content });
+    }
+  }
+
+  #emitReasoning(events: CompletionEvent[], text: string): void {
+    if (!text) return;
+    this.#reasoning += text;
+    events.push({ type: "reasoning", text });
+  }
+}
+
+export function createTimedFlowControl(
+  enabled: boolean,
+  intervalMs = 25,
+  now: () => number = () => performance.now(),
+  yieldTask: () => Promise<void> = () => new Promise<void>((resolve) => setImmediate(resolve)),
+): (emitted: boolean) => CompletionFlowControl {
+  let lastYield = now();
+  return (emitted) => {
+    if (!enabled || !emitted) return undefined;
+    const current = now();
+    if (current - lastYield < intervalMs) return undefined;
+    lastYield = current;
+    return yieldTask();
+  };
+}

package/src/serve/discovery-routes.ts

@@ -0,0 +1,205 @@
+import pkgJson from "../../package.json" with { type: "json" };
+import type { ServerContext } from "../server";
+import type { GenerationGateway } from "./generation-gateway";
+import { fit } from "../fit";
+import { isEmbeddingModel } from "../embed";
+import { Glm52Model } from "../model/glm52";
+
+const pkgVersion = (pkgJson as { version: string }).version;
+
+export type DiscoveryRoute =
+  | "library"
+  | "downloads"
+  | "api-index"
+  | "health"
+  | "models";
+
+export function matchDiscoveryRoute(method: string, pathname: string): DiscoveryRoute | null {
+  if (method !== "GET") return null;
+  switch (pathname) {
+    case "/library": return "library";
+    case "/downloads": return "downloads";
+    case "/v1": return "api-index";
+    case "/health": return "health";
+    case "/v1/models": return "models";
+    default: return pathname.startsWith("/v1/models/") ? "models" : null;
+  }
+}
+
+export interface DiscoveryRoutes {
+  handle(url: URL, request: Request): Promise<Response | null>;
+  invalidateLibrary(): void;
+}
+
+export function createDiscoveryRoutes(
+  ctx: ServerContext,
+  gateway: Pick<GenerationGateway, "batchMode">,
+  startedAt: number,
+): DiscoveryRoutes {
+  let libraryCache: { at: number; rows: unknown[] } | null = null;
+
+  return {
+    invalidateLibrary() {
+      libraryCache = null;
+    },
+
+    async handle(url, request) {
+      switch (matchDiscoveryRoute(request.method, url.pathname)) {
+        case "library": {
+          if (!libraryCache || Date.now() - libraryCache.at > 30_000) {
+            const { Registry, visionCapable, audioCapable } = await import("../registry");
+            const { loadModelConfig } = await import("../config");
+            const { supportTier } = await import("../model/support");
+            const registry = new Registry();
+            await registry.scan();
+            const rows = [];
+            for (const model of registry.listCanonical()) {
+              const tier = supportTier(model.modelType, model.repoId);
+              const supported = tier !== null;
+              let assessment = null;
+              try {
+                const config = await loadModelConfig(model.path);
+                const result = fit(
+                  config,
+                  model.sizeBytes,
+                  8192,
+                  undefined,
+                  undefined,
+                  model.expertsBytes,
+                );
+                assessment = {
+                  fits: result.fits,
+                  max_safe_context: result.maxSafeContext,
+                  predicted_decode_tps: result.predictedDecodeTps,
+                };
+              } catch {}
+              rows.push({
+                repo_id: model.repoId,
+                model_type: model.modelType,
+                size_bytes: model.sizeBytes,
+                quant_bits: model.quantBits,
+                vision: visionCapable(model),
+                audio: audioCapable(model),
+                supported,
+                support_tier: tier,
+                serving: model.repoId === ctx.modelId,
+                assessment,
+              });
+            }
+            libraryCache = { at: Date.now(), rows };
+          }
+          return Response.json({ models: libraryCache.rows });
+        }
+
+        case "downloads": {
+          const { downloadsSnapshot } = await import("../download");
+          return Response.json({ downloads: downloadsSnapshot() });
+        }
+
+        case "api-index":
+          return Response.json({
+            name: "mlx-bun",
+            version: pkgVersion,
+            model: ctx.modelId,
+            endpoints: [
+              "POST /v1/chat/completions",
+              "POST /v1/completions",
+              "POST /v1/messages",
+              "POST /v1/responses",
+              "POST /v1/embeddings",
+              "GET /v1/models",
+              "GET/POST/DELETE /v1/adapters",
+              "GET /health",
+              "GET /stats",
+              "GET /fit",
+              "GET /library",
+              "GET /downloads",
+            ],
+          });
+
+        case "health":
+          return new Response('{"status": "ok"}', {
+            headers: { "content-type": "application/json" },
+          });
+
+        case "models": {
+          const filterId = url.pathname.length > "/v1/models/".length - 1
+            ? decodeURIComponent(url.pathname.slice("/v1/models/".length))
+            : null;
+          const created = Math.floor(startedAt / 1000);
+          const genDefaults = {
+            temperature: ctx.genDefaults.temperature ?? null,
+            top_p: ctx.genDefaults.topP ?? null,
+            top_k: ctx.genDefaults.topK ?? null,
+          };
+          const isGlm52 = ctx.model instanceof Glm52Model;
+          const data: Array<Record<string, unknown>> = [{
+            id: ctx.modelId,
+            object: "model",
+            created,
+            owned_by: "mlx-bun",
+            context_window:
+              ctx.glmMemoryPlan?.contextTokens ??
+              ctx.model.config.text.maxPositionEmbeddings,
+            reasoning: ctx.template.supportsThinking,
+            vision: !!(ctx.vision || ctx.loadVision),
+            audio: !!(ctx.audio || ctx.loadAudio),
+            batch_mode: gateway.batchMode,
+            tools: true,
+            structured_output: true,
+            embeddings: isEmbeddingModel(ctx.model),
+            adapters: !isGlm52,
+            training: !isGlm52,
+            dsa: ctx.model instanceof Glm52Model && ctx.model.capabilities.dsa,
+            mtp: ctx.draft?.provider.id === "glm52-native-mtp",
+            capabilities: {
+              chat_completions: true,
+              text_completions: true,
+              anthropic_messages: true,
+              responses: true,
+              streaming: true,
+              tools: true,
+              structured_output: true,
+              logprobs: true,
+              embeddings: isEmbeddingModel(ctx.model),
+              vision: !!(ctx.vision || ctx.loadVision),
+              audio: !!(ctx.audio || ctx.loadAudio),
+              adapters: !isGlm52,
+              training: !isGlm52,
+            },
+            gen_defaults: genDefaults,
+          }];
+          try {
+            const { Registry, visionCapable } = await import("../registry");
+            const { supportTier } = await import("../model/support");
+            const registry = new Registry();
+            try {
+              if (registry.list().length === 0) await registry.scan();
+              for (const model of registry.listCanonical()) {
+                if (model.repoId === ctx.modelId) continue;
+                const tier = supportTier(model.modelType, model.repoId);
+                if (tier === null) continue;
+                data.push({
+                  id: model.repoId,
+                  object: "model",
+                  created,
+                  vision: visionCapable(model),
+                  tier,
+                });
+              }
+            } finally {
+              registry.close();
+            }
+          } catch {}
+          return Response.json({
+            object: "list",
+            data: filterId ? data.filter((model) => model.id === filterId) : data,
+          });
+        }
+
+        default:
+          return null;
+      }
+    },
+  };
+}