@atlaskit/editor-plugin-autocomplete 3.1.0 → 3.2.0

package/dist/es2019/pm-plugins/local-slow-lane-client.js

@@ -1,18 +1,28 @@
+import _defineProperty from "@babel/runtime/helpers/defineProperty";
 /**
  * Local Slow Lane Client: On-device inference via @mlc-ai/web-llm.
  *
- * Drop-in replacement for the network-based slow-lane-client.  Instead of
- * calling a backend API, this client uses MLC WebLLM to run a small language
- * model (SmolLM 135M) directly in the browser via WebGPU.
+ * Drop-in replacement for the network-based slow-lane-client. Instead of calling
+ * a backend API, this client runs two models in the browser via WebGPU, in a
+ * single MLCEngine, to reproduce the BE encoder's outputs on-device:
+ *
+ *   - Causal LM (SmolLM2-135M-Instruct): one decode step per word boundary. A
+ *     registered LogitProcessor captures the raw next-token logits, which
+ *     `computeBePayload` turns into a whole-word `lm_logits` payload — a faithful
+ *     port of the BE `CausalLMEncoder._get_top_k_probs` (masked softmax over the
+ *     vocab's first-tokens, prefix expansion, L2 reservation, log-space pooling).
+ *   - Semantic embedder (Snowflake Arctic Embed S): produces the real 384-d
+ *     `semantic_vector`. Inputs are wrapped as passages (see `wrapForArctic`) so
+ *     the runtime vector lands in the same space as the precomputed word bin.
  *
  * ── Why main thread (no Web Worker)? ─────────────────────────────────────
- * SmolLM 135M is small enough (~270 MB weights, 350-400 MB VRAM) that
- * WebGPU inference on the main thread is production-viable:
+ * The models are small enough (~640 MB combined VRAM) that WebGPU inference on
+ * the main thread is viable:
  *
  *   - WebGPU GPU compute is inherently async (doesn't block the main thread)
- *   - CPU overhead (tokenization + post-processing) is only 5-10 ms
- *   - Single forward pass latency is 50-150 ms — well within autocomplete
- *     expectations (~250 ms between word boundaries)
+ *   - CPU overhead (BE-parity post-processing) is a few ms
+ *   - Per-inference latency is well within autocomplete expectations
+ *     (~250 ms between word boundaries)
  *
  * This avoids all the complexity of Web Workers:
  *   - No CSP workarounds (blob URLs, inline scripts)
@@ -36,17 +46,380 @@ import { isWordBoundary } from './slow-lane-client';
 // ─── Constants ───────────────────────────────────────────────────────────────
 
 const DEFAULT_DEBOUNCE_MS = 300;
-export const LOCAL_MLC_MODEL_ID = 'SmolLM2-135M-Instruct-q0f16-MLC';
+export const LOCAL_MLC_CAUSAL_MODEL_ID = 'SmolLM2-135M-Instruct-q0f16-MLC';
+
+/**
+ * MLC ID for the semantic embedder (Snowflake Arctic Embed S, batch=4 variant).
+ *
+ * The `-b4` suffix selects the prebuilt variant compiled for a max batch size of
+ * 4 (≈239 MB VRAM) rather than `-b32` (≈1023 MB VRAM). Autocomplete embeds one
+ * context at a time, so `-b4` is the right fit. This model IS in
+ * `prebuiltAppConfig.model_list` of web-llm 0.2.82 — no `customModelConfig` needed.
+ */
+export const LOCAL_MLC_EMBEDDING_MODEL_ID = 'snowflake-arctic-embed-s-q0f32-MLC-b4';
+
+/**
+ * Wrap raw context text with BERT special tokens before embedding.
+ *
+ * web-llm's `EmbeddingPipeline` does NOT auto-prepend `[CLS]` / append `[SEP]`
+ * (the official MLC embeddings example wraps manually). The Python
+ * `sentence_transformers` side that generated the word-vector bin adds these
+ * inside `model.encode()`, so we must mirror it here for the runtime context
+ * vector to land in the same region of Arctic's embedding space as the bin.
+ *
+ * No query prefix is applied: the semantic step is sentence-to-sentence (`s2s`)
+ * similarity ("which words are conceptually similar to this context?"), not
+ * sentence-to-passage (`s2p`) retrieval. Arctic's query prefix would misframe
+ * the relationship. Encode both sides as passages. See implementation.md §4.3.
+ */
+export const wrapForArctic = text => `[CLS] ${text} [SEP]`;
+
+/**
+ * BE-parity constants — must match `CausalLMEncoder` defaults in the Python
+ * sidecar (`cc-smarts/python-sidecar/src/causal_lm_encoder.py`) and
+ * `SlowLaneEngine` (`typeahead_context_encoding.py`) so local payloads behave
+ * identically to the server-client setup.
+ */
+export const BE_PARITY = {
+  /** Final payload size cap (BE: `top_k_words`). */
+  TOP_K_WORDS: 2000,
+  /** L2 (domain) words admitted unconditionally before pooling (BE: `reserved_l2_slots`). */
+  RESERVED_L2_SLOTS: 500,
+  /** Log-space additive bias favouring L2 over L3 in the pool (BE: `l2_bias`). */
+  L2_BIAS: 1.0,
+  /** Drop words below this probability from the final payload (BE: `> 0.00001`). */
+  MIN_PROB: 0.00001,
+  /**
+   * Word-level approximation of the BE causal LM token limit.
+   *
+   * BE: `CausalLMEncoder.max_context_tokens = 100` (BPE tokens, left-truncated).
+   * FE: no tokenizer available, so we approximate with word count. English text
+   * averages ~1.3–1.5 BPE tokens/word, meaning 100 words ≈ 130–150 tokens.
+   * Using 100 words keeps the approximation simple and errs on the side of
+   * sending slightly more context than the BE sees — acceptable for a PoC.
+   */
+  MAX_CONTEXT_TOKENS: 100,
+  /**
+   * Word-level rolling window for the semantic embedder.
+   *
+   * BE: `SlowLaneEngine.max_context_words = 100` (applied in
+   * `typeahead_context_encoding.py` before calling `SemanticEncoder.encode`).
+   * Truncated identically here so the runtime Arctic vector lands in the same
+   * region of the embedding space as the precomputed word-vector bin.
+   */
+  MAX_CONTEXT_WORDS: 100
+};
+
+/**
+ * Return the last `n` whitespace-separated words of `text`, joined by spaces.
+ * Mirrors the BE rolling-window truncation applied before both encoders.
+ */
+const truncateToLastNWords = (text, n) => {
+  const words = text.trim().split(/\s+/u);
+  return words.length <= n ? text : words.slice(-n).join(' ');
+};
+
+// ─── Logit capture ─────────────────────────────────────────────────────────
+
+/**
+ * A LogitProcessor that captures the raw next-token logits and passes them
+ * through unmodified.
+ *
+ * web-llm invokes `processLogits` on the CPU after the model's forward pass and
+ * before sampling, handing us the full `Float32Array(vocab_size)` at the current
+ * decode position. We copy it off web-llm's shared buffer (which it may reuse
+ * across calls) and return the original untouched so sampling is unaffected.
+ *
+ * This is the raw-logit access the BE-parity algorithm needs (masked softmax +
+ * prefix expansion, consumed in a later step). Registered for the causal LM
+ * only — the embedder never decodes tokens, so it produces no logits.
+ */
+class CapturingLogitProcessor {
+  constructor() {
+    _defineProperty(this, "captured", null);
+    _defineProperty(this, "processLogits", logits => {
+      // Copy off web-llm's shared buffer — it may reuse `logits` across calls.
+      this.captured = new Float32Array(logits);
+      return logits;
+    });
+    _defineProperty(this, "processSampledToken", () => {
+      // No-op — we don't track sampled tokens.
+    });
+    _defineProperty(this, "resetState", () => {
+      this.captured = null;
+    });
+  }
+}
+
+// ─── BE-parity data + algorithm ──────────────────────────────────────────────
+
+/**
+ * Prefix-expansion map: first-token id → words whose space-prefixed SmolLM2
+ * encoding starts with that token. Generated offline by
+ * `scripts/gen_first_token_to_words.py`, which mirrors the BE's in-memory map
+ * (`CausalLMEncoder._ensure_loaded`).
+ *
+ * Populated lazily by `loadBePayloadData()` from a dynamically-imported JSON so
+ * the (large) payload is only fetched when the local client is actually
+ * initialised — keeping it out of the editor's main chunk for the vast majority
+ * of users (who run with `useLocalModel` off).
+ */
+let firstTokenToWords = new Map();
+
+/**
+ * L2 (Atlassian-domain) word set, derived from the keys of `vocabulary_10k.json`.
+ * Used by `computeBePayload` for tier-aware ranking: any word in the prefix map
+ * that is not in this set is treated as L3 (general English), matching the BE.
+ * Populated lazily alongside `firstTokenToWords` — see `loadBePayloadData()`.
+ */
+let l2Words = new Set();
+
+/**
+ * Array of token IDs that appear as a first token for at least one vocabulary
+ * word. Derived from `firstTokenToWords` when the data loads so `computeBePayload`
+ * does not re-allocate this array on every word-boundary call.
+ */
+let prefixMapTokenIds = [];
+
+/** De-dupes concurrent loads and lets repeated calls await the same payload. */
+let bePayloadDataPromise;
+
+/**
+ * Unwrap a dynamically imported JSON module to the parsed JSON value, working
+ * across the two interop modes AFM's bundler chain emits:
+ *
+ *   1. **`.default`-wrapped namespace** — classic webpack (and Jest) hang the
+ *      JSON value under the `default` export.
+ *   2. **Named-exports namespace** — webpack 5 / atlaspack with JSON
+ *      named-exports (or native ESM JSON modules) expose each top-level key as
+ *      a named export and shadow `default`, so `mod.default` can be `undefined`
+ *      (or some unrelated value) even though `mod` itself holds the data.
+ *
+ * The caller MUST declare the underlying JSON shape via `shape` because, in
+ * named-exports mode, a dense array `["a","b"]` and a sparse numeric-keyed
+ * object `{"5":"a","12":"b"}` are emitted identically (`{"0":..}` / `{"5":..}`);
+ * no runtime heuristic can tell them apart, so only the caller knows which:
+ *
+ *   - `'object'` — the JSON is a `{...}` (including sparse maps keyed by integer
+ *     IDs). The named exports are rebuilt into a plain object so `Object.entries`
+ *     yields the real keys, not synthetic array indices.
+ *   - `'array'` — the JSON is a `[...]`, reconstructed from the `0..n-1` indices.
+ *
+ * :param mod: The raw module object returned by `await import('./*.json')`.
+ * :param shape: `'object'` if the source JSON is `{...}`, `'array'` if `[...]`.
+ * :returns: The parsed JSON value, or `null` if neither interop mode applies.
+ */
+const unwrapJsonModule = (mod, shape) => {
+  if (mod == null || typeof mod !== 'object') {
+    return null;
+  }
+  const namespace = mod;
+
+  // Compute the named-export own-keys (strip synthetic markers).
+  const ownKeys = Object.keys(namespace).filter(k => k !== 'default' && k !== '__esModule');
+
+  // PREFER named exports when present — they always reflect the JSON's real
+  // top-level keys / indices, regardless of what `default` happens to be.
+  // Under JSON named-exports mode `default` is not necessarily the parsed
+  // value (e.g. for `{"service": 0, ...}` it can be the number `0`, with the
+  // real data in the named exports), so taking `default` first would corrupt it.
+  if (ownKeys.length > 0) {
+    if (shape === 'array') {
+      // JSON arrays are dense; reconstruct from `0..length-1` indices.
+      const len = ownKeys.length;
+      const arr = new Array(len);
+      for (let i = 0; i < len; i++) {
+        arr[i] = namespace[String(i)];
+      }
+      return arr;
+    }
+    // shape === 'object'. Rebuild a plain object from the (stripped) own
+    // keys so callers can `Object.entries()` it without iterating over
+    // `default` / `__esModule`, and to detach from the module-namespace
+    // object (which is sealed/non-extensible on some bundler outputs).
+    const obj = {};
+    for (const k of ownKeys) {
+      obj[k] = namespace[k];
+    }
+    return obj;
+  }
+
+  // Fallback: no named exports — classic webpack JSON-module interop where
+  // the whole parsed JSON value is hung under `default`. Trust it.
+  if ('default' in namespace && namespace.default != null) {
+    return namespace.default;
+  }
+  return null;
+};
 
-/** HF root for the default weights (includes `tensor-cache.json` for WebLLM 0.2+). */
-export const LOCAL_MLC_HF_MODEL_REPO = 'https://huggingface.co/mlc-ai/SmolLM2-135M-Instruct-q0f16-MLC';
-export const LOCAL_MLC_MODEL_LIB_WASM_NAME = 'SmolLM2-135M-Instruct-q0f16-ctx4k_cs1k-webgpu.wasm';
+/**
+ * Lazily load and build the BE-parity lookup tables from their JSON payloads.
+ * The dynamic imports are split into their own async chunks so neither file is
+ * bundled into the editor's main chunk unless local inference is initialised.
+ *
+ * :returns:
+ *   A promise that resolves once `firstTokenToWords`, `l2Words` and
+ *   `prefixMapTokenIds` are populated.
+ */
+const loadBePayloadData = () => {
+  if (!bePayloadDataPromise) {
+    bePayloadDataPromise = (async () => {
+      const [firstTokenToWordsModule, vocabularyModule] = await Promise.all([import( /* webpackChunkName: "@atlaskit-internal_editor-plugin-autocomplete-first-token-to-words" */'./data/first_token_to_words.json'), import( /* webpackChunkName: "@atlaskit-internal_editor-plugin-autocomplete-vocabulary-10k" */'./data/vocabulary_10k.json')]);
+      const firstTokenToWordsData = unwrapJsonModule(firstTokenToWordsModule, 'object');
+      const vocabularyData = unwrapJsonModule(vocabularyModule, 'object');
+      if (firstTokenToWordsData == null || (vocabularyData === null || vocabularyData === void 0 ? void 0 : vocabularyData.words) == null) {
+        // Hard-fail with a precise message so the catch() in initEngine logs
+        // exactly which import couldn't be unwrapped, rather than the generic
+        // V8 "Cannot convert undefined or null to object" we hit before the
+        // helper was added.
+        throw new Error(`[LocalSlowLane] JSON module could not be unwrapped — ` + `firstTokenToWordsData=${firstTokenToWordsData == null ? 'null/undefined' : 'defined'}, ` + `vocabularyData=${vocabularyData == null ? 'null/undefined' : vocabularyData.words == null ? 'defined but missing .words' : 'defined'}`);
+      }
+      firstTokenToWords = new Map(Object.entries(firstTokenToWordsData).map(([tokenId, words]) => [Number(tokenId), words]));
+      l2Words = new Set(Object.keys(vocabularyData.words));
+      prefixMapTokenIds = Array.from(firstTokenToWords.keys());
+      if (isAutocompleteDebugEnabled()) {
+        // eslint-disable-next-line no-console
+        console.log('%c[LocalSlowLane] %c✅ BE-parity payload data loaded:', 'color: #9c27b0; font-weight: bold;', 'color: #4caf50; font-weight: bold;', {
+          firstTokenToWordsEntries: firstTokenToWords.size,
+          l2WordsCount: l2Words.size,
+          prefixMapTokenIdsLength: prefixMapTokenIds.length
+        });
+      }
+    })().catch(e => {
+      // Don't cache a rejected promise — a transient import failure would
+      // otherwise prevent the local model from ever initialising again this
+      // session. Reset so the next init attempt retries.
+      bePayloadDataPromise = undefined;
+      throw e;
+    });
+  }
+  return bePayloadDataPromise;
+};
 
 /**
- * Original target repo (add-basics fine-tune). **Not compatible with WebLLM 0.2.x** (no `tensor-cache.json`).
- * @see module doc above
+ * Convert a raw next-token logit vector into a whole-word probability payload,
+ * faithfully porting the BE `CausalLMEncoder._get_top_k_probs`
+ * (`cc-smarts/python-sidecar/src/causal_lm_encoder.py`).
+ *
+ * Steps: (1) numerically-stable masked softmax over only the token ids present
+ * in the prefix-expansion map; (2) spread each token's probability to every
+ * whole word sharing that first token, taking the max; (3) reserve the top L2
+ * words unconditionally; (4) rank the remainder in a log-space pool with an
+ * additive L2 bias; (5) emit raw probabilities for the survivors, lowercased
+ * and trimmed at `MIN_PROB`.
+ *
+ * :params:
+ *   rawLogits: Full-vocabulary logits from the LM's single decode step
+ *   prefixMap: Map of first-token id to the words starting with that token
+ *   domainWords: Set of L2 (domain) words, for tier-aware ranking
+ * :returns:
+ *   A record of lowercase word to probability — the BE `lm_logits` payload
+ */
+export const computeBePayload = (rawLogits, prefixMap, domainWords,
+/**
+ * Pre-derived token-ID array for the softmax mask. Defaults to the
+ * module-level `prefixMapTokenIds` (zero allocation in production). Pass
+ * `Array.from(prefixMap.keys())` in tests that supply a custom prefixMap so
+ * the softmax mask stays consistent with the iteration in Step 2.
  */
-export const HUGGINGFACE_TB_SMOLLM_ADD_BASICS_REPO = 'https://huggingface.co/HuggingFaceTB/smollm-135M-instruct-add-basics-q0f16-MLC';
+validTokenIds = prefixMapTokenIds) => {
+  // 1. Numerically-stable masked softmax over validTokenIds only.
+  let maxLogit = -Infinity;
+  for (const id of validTokenIds) {
+    const v = rawLogits[id];
+    if (v > maxLogit) {
+      maxLogit = v;
+    }
+  }
+  let sumExp = 0;
+  const expByToken = new Map();
+  for (const id of validTokenIds) {
+    const e = Math.exp(rawLogits[id] - maxLogit);
+    expByToken.set(id, e);
+    sumExp += e;
+  }
+
+  // 2. Prefix expansion with max aggregation (probabilities sum to 1 over the
+  //    masked subset, so divide each token's exp by sumExp on the fly).
+  const wordProbs = new Map();
+  for (const [id, words] of prefixMap) {
+    var _expByToken$get;
+    const p = sumExp > 0 ? ((_expByToken$get = expByToken.get(id)) !== null && _expByToken$get !== void 0 ? _expByToken$get : 0) / sumExp : 0;
+    for (const w of words) {
+      var _wordProbs$get;
+      const prev = (_wordProbs$get = wordProbs.get(w)) !== null && _wordProbs$get !== void 0 ? _wordProbs$get : 0;
+      if (p > prev) {
+        wordProbs.set(w, p);
+      }
+    }
+  }
+
+  // 3. Split into L2 / L3 and reserve the top L2 slots unconditionally.
+  const l2Matches = [];
+  const l3Matches = [];
+  for (const [w, p] of wordProbs) {
+    if (domainWords.has(w)) {
+      l2Matches.push([w, p]);
+    } else {
+      l3Matches.push([w, p]);
+    }
+  }
+  l2Matches.sort((a, b) => b[1] - a[1]);
+  const reserved = l2Matches.slice(0, BE_PARITY.RESERVED_L2_SLOTS);
+
+  // 4. Pool the leftovers in log space; the L2 bias only affects ranking here.
+  // Words in l2Matches are unique and the array is sorted descending, so the
+  // non-reserved entries are exactly the tail after the reserved prefix — slice
+  // it directly rather than allocating a Set and scanning every entry on this
+  // hot path (runs ~every word boundary while typing).
+  const pool = [];
+  for (const [w, p] of l2Matches.slice(BE_PARITY.RESERVED_L2_SLOTS)) {
+    pool.push([w, Math.log(Math.max(p, 1e-10)) + BE_PARITY.L2_BIAS]);
+  }
+  for (const [w, p] of l3Matches) {
+    pool.push([w, Math.log(Math.max(p, 1e-10))]);
+  }
+  pool.sort((a, b) => b[1] - a[1]);
+  const remainingSlots = Math.max(0, BE_PARITY.TOP_K_WORDS - reserved.length);
+  const poolWinners = pool.slice(0, remainingSlots);
+
+  // 5. Assemble payload: store RAW probabilities (the bias was ranking-only),
+  // lowercase keys, trimmed at MIN_PROB. Reserved first, then pool winners.
+  // Reserved entries are written first; pool-winner writes must NOT clobber a
+  // reserved entry whose normalised key collides (two source words can
+  // `.trim().toLowerCase()` to the same key — e.g. "Function" vs "function ").
+  // Without the existence guard, a low-probability pool winner would silently
+  // overwrite the (higher-probability) reserved entry, degrading top-K
+  // quality in a way that's invisible from the debug summary.
+  const result = {};
+  const addEntry = (word, prob, allowOverwrite) => {
+    if (prob <= BE_PARITY.MIN_PROB) {
+      return;
+    }
+    const key = word.trim().toLowerCase();
+    if (!allowOverwrite && key in result) {
+      return;
+    }
+    result[key] = prob;
+  };
+  for (const [w, p] of reserved) {
+    addEntry(w, p, true);
+  }
+  for (const [w] of poolWinners) {
+    var _wordProbs$get2;
+    addEntry(w, (_wordProbs$get2 = wordProbs.get(w)) !== null && _wordProbs$get2 !== void 0 ? _wordProbs$get2 : 0, false);
+  }
+  if (isAutocompleteDebugEnabled()) {
+    const topReserved = reserved.slice(0, 5).map(([w, p]) => `${w}:${(p * 100).toFixed(2)}%`).join(', ');
+    const topPool = poolWinners.slice(0, 5).map(([w]) => {
+      var _wordProbs$get3;
+      return `${w}:${(((_wordProbs$get3 = wordProbs.get(w)) !== null && _wordProbs$get3 !== void 0 ? _wordProbs$get3 : 0) * 100).toFixed(2)}%`;
+    }).join(', ');
+    // eslint-disable-next-line no-console
+    console.log('%c[computeBePayload] %c%d valid tokens → %d words expanded | L2: %d / L3: %d | reserved: %d | pool winners: %d | final: %d words\n  maxLogit(masked): %s | sumExp: %s\n  top reserved L2: %s\n  top pool: %s', 'color: #9c27b0; font-weight: bold;', 'color: inherit;', validTokenIds.length, wordProbs.size, l2Matches.length, l3Matches.length, reserved.length, poolWinners.length, Object.keys(result).length, maxLogit.toFixed(3), sumExp.toFixed(1), topReserved || '(none)', topPool || '(none)');
+  }
+  return result;
+};
 
 // ─── Factory ─────────────────────────────────────────────────────────────────
 
@@ -73,7 +446,7 @@ export const createLocalSlowLaneClient = (config = {}) => {
     debounceMs = DEFAULT_DEBOUNCE_MS,
     onUpdate,
     onStatus,
-    modelId = LOCAL_MLC_MODEL_ID,
+    modelId = LOCAL_MLC_CAUSAL_MODEL_ID,
     customModelConfig
   } = config;
 
@@ -89,6 +462,9 @@ export const createLocalSlowLaneClient = (config = {}) => {
   let initFailed = false;
   let engine = null;
   let engineInitPromise = null;
+  // Captures raw next-token logits from the LM's single decode step. Registered
+  // with the engine below; `lmLogitsCapture.captured` is consumed in a later step.
+  const lmLogitsCapture = new CapturingLogitProcessor();
   const unloadEngine = engineToUnload => {
     engineToUnload.unload().catch(error => {
       if (isAutocompleteDebugEnabled()) {
@@ -112,16 +488,19 @@ export const createLocalSlowLaneClient = (config = {}) => {
     try {
       if (isAutocompleteDebugEnabled()) {
         // eslint-disable-next-line no-console
-        console.log(`%c[LocalSlowLane] %c🚀 Initialising MLC engine with model: ${modelId}`, 'color: #9c27b0; font-weight: bold;', 'color: inherit;');
+        console.log(`%c[LocalSlowLane] %c🚀 Initialising MLC engine with models: ${modelId} (LM) + ${LOCAL_MLC_EMBEDDING_MODEL_ID} (embedder)`, 'color: #9c27b0; font-weight: bold;', 'color: inherit;');
       }
-      onStatus === null || onStatus === void 0 ? void 0 : onStatus(`Initialising model: ${modelId}…`);
+      onStatus === null || onStatus === void 0 ? void 0 : onStatus(`Initialising models: ${modelId} + ${LOCAL_MLC_EMBEDDING_MODEL_ID}…`);
       if (!('gpu' in navigator)) {
         throw new Error('WebGPU not supported');
       }
-      const {
-        CreateMLCEngine,
+
+      // Fetch the web-llm runtime and the BE-parity lookup tables in parallel;
+      // both are dynamically imported so they stay out of the main editor chunk.
+      const [{
+        MLCEngine: MLCEngineCtor,
         prebuiltAppConfig
-      } = await import( /* webpackChunkName: "@atlaskit-internal_editor-plugin-autocomplete-mlc-web-llm" */'@mlc-ai/web-llm');
+      }] = await Promise.all([import( /* webpackChunkName: "@atlaskit-internal_editor-plugin-autocomplete-mlc-web-llm" */'@mlc-ai/web-llm'), loadBePayloadData()]);
       const customModelRecord = customModelConfig ? {
         model: customModelConfig.model,
         model_id: modelId,
@@ -140,20 +519,32 @@ export const createLocalSlowLaneClient = (config = {}) => {
       const appConfig = {
         model_list: [...prebuiltAppConfig.model_list, ...(customModelRecord ? [customModelRecord] : [])]
       };
-      engine = await CreateMLCEngine(modelId, {
+
+      // Construct the engine with the logit-capture processor registered for
+      // the causal LM only (the embedder never decodes tokens), then load
+      // both the LM and the embedder into the same engine (multi-model).
+      const newEngine = new MLCEngineCtor({
         appConfig,
-        initProgressCallback
+        initProgressCallback,
+        logitProcessorRegistry: new Map([[modelId, lmLogitsCapture]])
       });
+      await newEngine.reload([modelId, LOCAL_MLC_EMBEDDING_MODEL_ID]);
       if (destroyed) {
         // destroy() was called while we were loading — clean up
-        unloadEngine(engine);
-        engine = null;
+        unloadEngine(newEngine);
         return;
       }
+      engine = newEngine;
       ready = true;
       if (isAutocompleteDebugEnabled()) {
         // eslint-disable-next-line no-console
-        console.log('%c[LocalSlowLane] %c✅ MLC engine loaded and ready', 'color: #9c27b0; font-weight: bold;', 'color: #4caf50;');
+        console.log('%c[LocalSlowLane] %c✅ Both models loaded and ready', 'color: #9c27b0; font-weight: bold;', 'color: #4caf50;');
+        // One-time identity summary so you can confirm which models are active
+        // without digging through the init-progress scroll.
+        // eslint-disable-next-line no-console
+        console.log('%c[LocalSlowLane] %c🧠 Causal LM    →', 'color: #9c27b0; font-weight: bold;', 'color: #2196f3; font-weight: bold;', modelId);
+        // eslint-disable-next-line no-console
+        console.log('%c[LocalSlowLane] %c🔢 Embedder     →', 'color: #9c27b0; font-weight: bold;', 'color: #009688; font-weight: bold;', LOCAL_MLC_EMBEDDING_MODEL_ID);
       }
       onStatus === null || onStatus === void 0 ? void 0 : onStatus('Model loaded and ready.');
     } catch (err) {
@@ -181,85 +572,102 @@ export const createLocalSlowLaneClient = (config = {}) => {
   // ── Inference ──────────────────────────────────────────────────────────
 
   /**
-   * Run a single forward pass to extract next-token logit probabilities.
+   * Run a single forward pass to produce the BE-parity slow-lane outputs.
    *
-   * We use the chat completions API with `max_tokens: 1` and `logprobs: true`
-   * to get the model's next-token distribution without generating text.
-   * This is the cheapest possible inference call — a single forward pass.
+   * Two calls run in parallel on the shared engine:
+   *   - `completions.create({ max_tokens: 1 })` runs the causal LM for exactly
+   *     one decode step. We ignore the generated text; the LogitProcessor
+   *     captures the raw next-token logits during that step, which we turn into
+   *     a whole-word payload via `computeBePayload`.
+   *   - `embeddings.create(...)` runs the Arctic embedder to produce the real
+   *     384-d semantic vector (passage-encoded; see `wrapForArctic`).
    */
   const runInference = async (text, requestId) => {
     if (!engine || destroyed) {
       return;
     }
+
+    // Clear the capture buffer so we read only this pass's logits. The engine
+    // serialises per-model requests and updateContext is debounced, so the
+    // latest request's decode step is the last to populate `captured` before
+    // we read it below; stale requests bail on the latestRequestId guard.
+    lmLogitsCapture.resetState();
+
+    // Apply BE-parity rolling-window truncation before both encoders.
+    // BE semantic: last max_context_words words (typeahead_context_encoding.py:36)
+    // BE causal LM: last max_context_tokens BPE tokens (causal_lm_encoder.py:194–198),
+    //               approximated here with word count (no tokenizer available on FE).
+    const lmText = truncateToLastNWords(text, BE_PARITY.MAX_CONTEXT_TOKENS);
+    const semanticText = truncateToLastNWords(text, BE_PARITY.MAX_CONTEXT_WORDS);
+    const arcticInput = wrapForArctic(semanticText);
+    if (isAutocompleteDebugEnabled()) {
+      // eslint-disable-next-line no-console
+      console.log(`%c[LocalSlowLane] %c🔢 Arctic input (${arcticInput.length} chars, ${semanticText.split(/\s+/u).length} words): "${arcticInput.length > 100 ? `${arcticInput.slice(0, 100)}…` : arcticInput}"`, 'color: #9c27b0; font-weight: bold;', 'color: #009688;');
+      // eslint-disable-next-line no-console
+      console.log(`%c[LocalSlowLane] %c🧠 LM input (${lmText.length} chars, ${lmText.split(/\s+/u).length} words): "${lmText.length > 100 ? `${lmText.slice(0, 100)}…` : lmText}"`, 'color: #9c27b0; font-weight: bold;', 'color: #2196f3;');
+    }
     try {
-      var _response$choices, _response$choices$, _response$choices$$lo;
-      // Use chat completion with logprobs to get next-token distribution
-      const response = await engine.chat.completions.create({
-        messages: [{
-          role: 'user',
-          content: text
-        }],
+      var _embeddingResponse$da, _embeddingResponse$da2;
+      const tStart = performance.now();
+      let tLmDone = 0;
+      let tEmbDone = 0;
+      const [, embeddingResponse] = await Promise.all([engine.completions.create({
+        model: modelId,
+        prompt: lmText,
         max_tokens: 1,
-        logprobs: true,
-        top_logprobs: 5,
-        temperature: 0
-      });
+        temperature: 0,
+        logprobs: false
+      }).then(r => {
+        tLmDone = performance.now();
+        return r;
+      }), engine.embeddings.create({
+        model: LOCAL_MLC_EMBEDDING_MODEL_ID,
+        input: arcticInput
+      }).then(r => {
+        tEmbDone = performance.now();
+        return r;
+      })]);
+      if (isAutocompleteDebugEnabled()) {
+        // eslint-disable-next-line no-console
+        console.log(`%c[LocalSlowLane] %c⏱ LM: ${(tLmDone - tStart).toFixed(0)}ms | Embedder: ${(tEmbDone - tStart).toFixed(0)}ms | Total: ${(Math.max(tLmDone, tEmbDone) - tStart).toFixed(0)}ms`, 'color: #9c27b0; font-weight: bold;', 'color: #ff9800;');
+      }
 
       // Discard stale results
       if (requestId < latestRequestId || destroyed) {
         return;
       }
 
-      // ── Extract LM logits ───────────────────────────────────────
-      const lmLogits = {};
-      const logprobsContent = (_response$choices = response.choices) === null || _response$choices === void 0 ? void 0 : (_response$choices$ = _response$choices[0]) === null || _response$choices$ === void 0 ? void 0 : (_response$choices$$lo = _response$choices$.logprobs) === null || _response$choices$$lo === void 0 ? void 0 : _response$choices$$lo.content;
-      if (logprobsContent && logprobsContent.length > 0) {
-        const tokenLogprobs = logprobsContent[0];
-
-        // Add the top token
-        if (tokenLogprobs.token) {
-          const token = tokenLogprobs.token.trim().toLowerCase();
-          // @ts-ignore TS1501: Unicode regex flag requires a newer TS target than the declaration build uses.
-          if (token.length > 0 && /^[a-z]/iu.test(token)) {
-            lmLogits[token] = Math.exp(tokenLogprobs.logprob);
-          }
-        }
-
-        // Add alternative tokens from top_logprobs
-        if (tokenLogprobs.top_logprobs) {
-          for (const alt of tokenLogprobs.top_logprobs) {
-            const token = alt.token.trim().toLowerCase();
-            // @ts-ignore TS1501: Unicode regex flag requires a newer TS target than the declaration build uses.
-            if (token.length > 0 && /^[a-z]/iu.test(token)) {
-              lmLogits[token] = Math.exp(alt.logprob);
-            }
-          }
-        }
-      }
-      storedLmLogits = Object.keys(lmLogits).length > 0 ? lmLogits : null;
-
-      // ── Semantic vector ─────────────────────────────────────────
-      // SmolLM is a generative model, not an embedding model, so we
-      // don't get a true semantic vector. We generate a lightweight
-      // pseudo-embedding from the logit distribution for compatibility
-      // with the existing scoring pipeline.
-      //
-      // For a production implementation, you would use a dedicated
-      // embedding model (e.g. via web-llm's embeddings API with an
-      // embedding-specific model).
-      if (storedLmLogits) {
-        const logitValues = Object.values(storedLmLogits);
-        storedContextVector = new Float32Array(logitValues);
+      // ── LM logits: whole-word BE-parity payload ──────────────────
+      const rawLogits = lmLogitsCapture.captured;
+      if (rawLogits) {
+        const payload = computeBePayload(rawLogits, firstTokenToWords, l2Words);
+        storedLmLogits = Object.keys(payload).length > 0 ? payload : null;
       } else {
-        storedContextVector = null;
+        storedLmLogits = null;
       }
+
+      // ── Semantic vector: real 384-d Arctic embedding ─────────────
+      // Guard against base64-encoded responses (encoding_format: 'base64' would
+      // yield a string, and new Float32Array(string) silently produces an empty
+      // array, corrupting downstream cosine-similarity scoring).
+      const embedding = (_embeddingResponse$da = embeddingResponse.data) === null || _embeddingResponse$da === void 0 ? void 0 : (_embeddingResponse$da2 = _embeddingResponse$da[0]) === null || _embeddingResponse$da2 === void 0 ? void 0 : _embeddingResponse$da2.embedding;
+      storedContextVector = Array.isArray(embedding) && embedding.length > 0 ? new Float32Array(embedding) : null;
       if (isAutocompleteDebugEnabled()) {
         // eslint-disable-next-line no-console
         console.groupCollapsed(`%c[LocalSlowLane] %c📥 Inference result (request #${requestId})`, 'color: #9c27b0; font-weight: bold;', 'color: inherit;');
+        if (storedContextVector) {
+          let sumSq = 0;
+          for (let i = 0; i < storedContextVector.length; i++) {
+            sumSq += storedContextVector[i] * storedContextVector[i];
+          }
+          // eslint-disable-next-line no-console
+          console.log(`✅ semantic vector: ${storedContextVector.length} dims (L2 norm ${Math.sqrt(sumSq).toFixed(3)})`);
+        } else {
+          // eslint-disable-next-line no-console
+          console.log('❌ No vector');
+        }
         // eslint-disable-next-line no-console
-        console.log(storedContextVector ? `✅ pseudo-vector: ${storedContextVector.length} dims` : '❌ No vector');
-        // eslint-disable-next-line no-console
-        console.log(storedLmLogits ? `✅ lm_logits: ${Object.keys(storedLmLogits).length} tokens` : '❌ No lm_logits');
+        console.log(storedLmLogits ? `✅ lm_logits: ${Object.keys(storedLmLogits).length} words` : '❌ No lm_logits');
         if (storedLmLogits) {
           const topTokens = Object.entries(storedLmLogits).sort(([, a], [, b]) => b - a).slice(0, 10);
           // eslint-disable-next-line no-console