mlx-bun 0.2.0 → 0.4.0

package/src/serve/session-completion-engine.ts

@@ -0,0 +1,110 @@
+import { cleanupFailure } from "../engine/resources";
+import type { ExecutionPlan, InferenceMethod, Timer } from "../contracts/generation";
+import { CancellationSource, throwIfCancelled } from "../engine/cancellation";
+import { createInferenceEngine } from "../engine/engine";
+import type { GenerateStats } from "../generate";
+import type { CompletionEngine } from "./completion-executor";
+
+const timer: Timer = {
+  after(ms, callback) {
+    const handle = setTimeout(callback, ms);
+    return () => clearTimeout(handle);
+  },
+};
+
+type RuntimeArguments = Parameters<CompletionEngine["run"]>;
+
+export interface CompletionRunRequest {
+  prompt: RuntimeArguments[0]; options: RuntimeArguments[1]; vision?: RuntimeArguments[3];
+  shape: RuntimeArguments[4]; placement: RuntimeArguments[5];
+  signal?: RuntimeArguments[6]; trace?: RuntimeArguments[7];
+}
+
+/** Bind one selected runtime operation to the method-neutral session. */
+export function createCompletionMethod(
+  runtime: CompletionEngine, request: CompletionRunRequest,
+  onStarted?: () => void, onFailure?: (error: unknown) => void,
+): InferenceMethod<GenerateStats> {
+  const { prompt, options, vision, shape, placement, signal, trace } = request;
+  const maxTokens = options.maxTokens ?? 512;
+  return {
+    id: placement.execution?.method ?? placement.mechanism,
+    async createRun(cancel) {
+      throwIfCancelled(cancel);
+      const controller = new AbortController();
+      const unsubscribe = cancel.subscribe(() => controller.abort(
+        signal?.reason ?? new DOMException("generation cancelled", "AbortError"),
+      ));
+      return {
+        async execute(output) {
+          onStarted?.();
+          try {
+            const stats = await runtime.run(prompt, options,
+              (token, logprobs) => output.commit([token], logprobs ? [logprobs] : undefined),
+              vision, shape, placement, controller.signal, trace);
+            return {
+              finishReason: stats.finishReason ?? (stats.generatedTokens >= maxTokens ? "length" : "stop"),
+              metrics: stats,
+            };
+          } catch (error) { onFailure?.(error); throw error; }
+        },
+        async close() { unsubscribe(); },
+      };
+    },
+  };
+}
+
+/** HTTP already bounds its output in CompletionSink. Deliver through the same
+ * session lifecycle without inserting another queue between sampling and the
+ * stop/tool parser. The runtime retains scheduling and native resource ownership. */
+export function createSessionCompletionEngine(
+  runtime: CompletionEngine,
+  releaseUnstarted: (options: Parameters<CompletionEngine["run"]>[1],
+    vision: Parameters<CompletionEngine["run"]>[3]) => void,
+): CompletionEngine & { close(): Promise<void> } {
+  const engine = createInferenceEngine<ExecutionPlan<GenerateStats>, GenerateStats>({
+    async plan(plan, cancellation) { throwIfCancelled(cancellation); return plan; },
+  }, { timer });
+  return {
+    place: (shape, options) => runtime.place(shape, options),
+    async run(prompt, options, onToken, vision, shape, placement, signal, trace) {
+      const cancellation = new CancellationSource();
+      const abort = () => cancellation.cancel("requested");
+      signal?.addEventListener("abort", abort, { once: true });
+      if (signal?.aborted) abort();
+      let started = false;
+      let failure: unknown;
+      let failed = false;
+      const maxTokens = options.maxTokens ?? 512;
+      const method = createCompletionMethod(runtime,
+        { prompt, options, vision, shape, placement, signal, trace },
+        () => { started = true; },
+        (error) => { failed = true; failure = error; });
+      try {
+        const session = await engine.open({ id: method.id, method, outputTokenLimit: maxTokens }, {
+          output: "callback", cancellation,
+          async onTokens(ids, logprobs) {
+            for (let i = 0; i < ids.length; i++)
+              if (await onToken(ids[i]!, logprobs?.[i]) === false) return false;
+          },
+        });
+        const outcome = await session.result;
+        if (outcome.status === "completed") return outcome.result.metrics;
+        if (failed) throw failure; // preserve native/HTTP error identity and stack
+        if (outcome.status === "cancelled")
+          throw signal?.reason ?? new DOMException("generation cancelled", "AbortError");
+        throw new Error(outcome.error.message);
+      } catch (error) {
+        if (!started) {
+          started = true; // release once, preserving the original execution error
+          cleanupFailure(error, () => releaseUnstarted(options, vision));
+        }
+        throw error;
+      } finally {
+        signal?.removeEventListener("abort", abort);
+        if (!started) releaseUnstarted(options, vision);
+      }
+    },
+    close: () => engine.close(),
+  };
+}

package/src/serve/session-search.ts

@@ -258,7 +258,7 @@ export async function searchSessions(sessionDir: string, query: string): Promise
    REST handlers — pure (URL) => Response, no `ctx` dependency, same
    convention as src/memory/rest.ts / src/hub-rest.ts (dispatched from
    server.ts by path+method; unit-tested directly in
-   tests/session-search.test.ts without booting a model context).
+   tests/serve/session-search.test.ts without booting a model context).
    ──────────────────────────────────────────────────────────────────── */
 
 function jsonOk<T extends object>(body: T, init?: ResponseInit): Response {

package/src/serve/static-routes.ts

@@ -55,7 +55,7 @@ export function handleStaticRoute(
     case "/curves": {
       let html = assets.curvePage;
       try {
-        html = readFileSync(new URL("../assets/curve-designer.html", import.meta.url), "utf8");
+        html = readFileSync(new URL("../lab/curve/curve-designer.html", import.meta.url), "utf8");
       } catch {
         // Compiled binaries use the embedded copy.
       }
@@ -66,7 +66,7 @@ export function handleStaticRoute(
     case "/dag":
       try {
         const html = readFileSync(
-          new URL("../../docs/dag/training-inference-map.html", import.meta.url),
+          new URL("./assets/training-inference-map.html", import.meta.url),
           "utf8",
         );
         return new Response(html, {
@@ -74,14 +74,14 @@ export function handleStaticRoute(
         });
       } catch {
         return new Response(
-          "DAG map artifact not found at docs/dag/training-inference-map.html",
+          "DAG map artifact not found",
           { status: 404 },
         );
       }
     case "/curve-terrain":
       try {
         const html = readFileSync(
-          new URL("../../docs/archive/investigations/curve-terrain.html", import.meta.url),
+          new URL("./assets/curve-terrain.html", import.meta.url),
           "utf8",
         );
         return new Response(html, {
@@ -89,7 +89,7 @@ export function handleStaticRoute(
         });
       } catch {
         return new Response(
-          "curve terrain artifact not found; run scripts/experiments/curve-terrain.ts first",
+          "curve terrain artifact not found",
           { status: 404 },
         );
       }

package/src/serve/text-completion-stage.ts

@@ -0,0 +1,104 @@
+// TextCompletionStage: TextCompletionRequest in, InferenceRequest out.
+// mlx_lm.server's raw text completion: NO chat template — the prompt string
+// is tokenized directly (tokenizer.encode with the tokenizer's own
+// special-token handling, exactly mlx-lm's `tokenizer.encode(request.prompt)`).
+// Same admission / adapter / lane path as chat; no tool router or thinking
+// splitter (raw text in, raw text out).
+import type { ServingContext } from "./model-host";
+import type { TextCompletionRequest } from "./chat-request";
+import type { InferenceRequest } from "./inference-request";
+import { RequestError } from "./pipeline";
+import { RequestOwnership } from "./request-plan";
+import type { RequestPrep } from "./request-prep";
+import { StopMatcher, ThinkingTagSplitter, ToolAwareStream } from "./token-streams";
+
+export class TextCompletionStage {
+  constructor(
+    private readonly ctx: ServingContext,
+    private readonly prep: RequestPrep,
+    private readonly contextLimit: number | null,
+    /** `--max-tokens` / the GLM plan's reservation. No invented fallback:
+     *  unset runs to EOS or the admitted context (with no template an EOS
+     *  may never come — that run is bounded by admission, the real limit). */
+    private readonly defaultGeneratedTokens: number | undefined,
+    private readonly defaultAdapter?: string,
+  ) {}
+
+  async run(request: TextCompletionRequest, requestId?: string): Promise<InferenceRequest> {
+    const { ctx, prep } = this;
+    const body = request.params;
+    const id = requestId ?? `cmpl-${crypto.randomUUID()}`;
+    const promptIds = ctx.tokenizer.encode(body.prompt);
+    const ownership = new RequestOwnership();
+    const warnings: string[] = [];
+    const reject = (message: string): never => {
+      ownership.dispose();
+      throw new RequestError(400, message);
+    };
+    let options: ReturnType<RequestPrep["toOptions"]>;
+    try {
+      options = prep.toOptions(body as never);
+    } catch (e) {
+      // bad logit_bias (non-numeric keys/values) — mlx-lm's coercion error
+      return reject((e as Error).message);
+    }
+    // Grammar-constrained decoding on raw completions too (response_format /
+    // guided_* / structured_outputs). The degrade path has no chat template
+    // to inject a system message into, so a degrade only emits the Warning
+    // header (no prompt injection) — documented gap vs the chat lane, which
+    // mirrors oMLX's text-completions behavior.
+    const grammarReq =
+      body.response_format != null || !!body.guided_grammar ||
+      !!body.guided_regex || !!body.guided_choice?.length ||
+      body.structured_outputs != null;
+    if (grammarReq) {
+      const g = await prep.compileGrammarForRequest(body as never);
+      if (g.controller) options.grammar = ownership.own(g.controller);
+      else if (g.degradeHint)
+        warnings.push(
+          `grammar not enforced: ${g.degradeHint} - no prompt injection on /v1/completions`);
+    }
+    // Unset = no request-level ceiling; admission clamps to the admitted
+    // context. DEVIATION from mlx_lm.server (which stops a defaulted
+    // request at 512): pass --max-tokens 512 to reproduce the reference.
+    const requestedMaxTokens = body.max_completion_tokens ?? body.max_tokens ??
+      this.defaultGeneratedTokens ?? Infinity;
+    let adapterIds: string[];
+    try {
+      adapterIds = ctx.adapters.resolveSpec(body.adapter ?? this.defaultAdapter);
+    } catch (e) {
+      return reject((e as Error).message);
+    }
+    const wantLogprobs = body.logprobs === true;
+    const topLogprobs =
+      typeof body.top_logprobs === "number" && body.top_logprobs > 0
+        ? body.top_logprobs : 0;
+    return {
+      requestId: id,
+      stream: request.stream,
+      warnings,
+      plan: {
+        promptIds,
+        options,
+        requestedMaxTokens,
+        contextLimit: this.contextLimit,
+        stream: request.stream,
+        wantLogprobs,
+        topLogprobs,
+        adapterIds,
+        hasVision: false,
+        userSeed: body.seed !== undefined,
+        hasGrammar: !!options.grammar,
+        hasDraft: !!ctx.draft,
+        ownership,
+      },
+      pipeline: {
+        router: new ToolAwareStream(ctx.tokenizer, "plain", null),
+        stopper: new StopMatcher(options.stopSequences),
+        thinking: new ThinkingTagSplitter(false),
+        collectToolCalls: false,
+      },
+      idToToken: (tokenId) => ctx.tokenizer.idToToken(tokenId),
+    };
+  }
+}