@atlaskit/editor-plugin-autocomplete 3.2.0 → 3.4.0

package/src/pm-plugins/autocomplete-plugin.ts

@@ -14,6 +14,7 @@ import type { EditorView } from '@atlaskit/editor-prosemirror/view';
 
 import type { AutocompletePlugin } from '../autocompletePluginType';
 
+import { isAutocompleteDebugEnabled } from './debug-mode';
 import { createGhostTextDecorationSet } from './ghost-text-decoration';
 import { createLocalSlowLaneClient, type LocalSlowLaneClient } from './local-slow-lane-client';
 import { createSlowLaneClient, setDefaultSlowLaneClient, isWordBoundary } from './slow-lane-client';
@@ -28,6 +29,19 @@ import {
 export const autocompletePluginKey: PluginKey = new PluginKey('autocomplete');
 
 const DEBOUNCE_MS = 150;
+const NETWORK_SLOW_LANE_DEBOUNCE_MS = 300;
+const LOCAL_SLOW_LANE_DEBOUNCE_MS = 100;
+const CONTEXT_REFRESH_THROTTLE_MS = 1000;
+// Caps the dedup Set so long editing sessions don't retain every distinct
+// version of the (potentially hundreds-of-KB) page content for the plugin's
+// lifetime. Eviction is FIFO; a re-ingest of an evicted text is harmless.
+const MAX_INGESTED_CONTEXT_TEXTS = 50;
+// Caps how many times the word-boundary path will retry getContext() while the
+// parent comment is still missing. Combined with the 1s throttle this gives a
+// ~5s window to cover a still-loading comment thread, then stops permanently so
+// non-comment editors (where parentCommentContent never arrives) don't refetch
+// on every word boundary for the plugin's lifetime.
+const MAX_CONTEXT_REFRESH_ATTEMPTS = 5;
 
 const hasDestroy = (
 	client: ReturnType<typeof createSlowLaneClient> | LocalSlowLaneClient,
@@ -241,6 +255,12 @@ export const createAutocompletePlugin = (
 	let debounceTimer: ReturnType<typeof setTimeout> | null = null;
 	let hasIngestedPage = false;
 	let resolvedContext: AutocompleteContext | undefined;
+	/**
+	 * Kept in sync with the live EditorView so the async getContext() promise
+	 * can re-trigger a slow-lane update the moment context arrives, even if the
+	 * user has already typed several words before the promise resolved.
+	 */
+	let currentView: EditorView | null = null;
 	/**
 	 * Set after accepting a suggestion so the next doc-change update
 	 * skips scheduling a new prediction for the just-inserted text.
@@ -285,14 +305,139 @@ export const createAutocompletePlugin = (
 
 	const slowLaneClient = options?.useLocalModel
 		? createLocalSlowLaneClient({
-				debounceMs: 300,
+				debounceMs: LOCAL_SLOW_LANE_DEBOUNCE_MS,
 			})
 		: createSlowLaneClient({
 				baseUrl: '',
-				debounceMs: 300,
+				debounceMs: NETWORK_SLOW_LANE_DEBOUNCE_MS,
 			});
 	setDefaultSlowLaneClient(slowLaneClient);
 
+	let contextRequestInFlight = false;
+	let lastContextRefreshAt = 0;
+	// Bounds the word-boundary retry loop so it terminates even when the editor is
+	// not in a comment thread (parentCommentContent never resolves).
+	let wordBoundaryRefreshAttempts = 0;
+	// Set when the plugin is torn down so in-flight getContext() resolutions don't
+	// mutate the global text-predictor state after destruction.
+	let destroyed = false;
+	const ingestedContextTexts = new Set<string>();
+
+	const logContextResolved = (source: string, context?: AutocompleteContext): void => {
+		if (!isAutocompleteDebugEnabled()) {
+			return;
+		}
+
+		// eslint-disable-next-line no-console
+		console.log(
+			'%c[Autocomplete] %cgetContext resolved',
+			'color: #00b8d9; font-weight: bold;',
+			'color: inherit;',
+			{
+				source,
+				hasParentComment: !!context?.parentCommentContent,
+				parentCommentPreview: context?.parentCommentContent?.slice(0, 80),
+				siblingCount: context?.siblingCommentsContents?.length ?? 0,
+				hasFullPage: !!context?.fullPageContent,
+			},
+		);
+	};
+
+	const applyContext = (context: AutocompleteContext): void => {
+		// Merge rather than replace: the word-boundary retry may resolve only a
+		// late-arriving field (e.g. parentCommentContent) without re-sending
+		// fullPageContent, so replacing would drop previously resolved context.
+		// Strip undefined values first so a field explicitly set to undefined by
+		// getContext doesn't overwrite a previously resolved value.
+		const definedContext = Object.fromEntries(
+			Object.entries(context).filter(([, value]) => value !== undefined),
+		);
+		resolvedContext = { ...resolvedContext, ...definedContext };
+
+		const ingestContextText = (text?: string): void => {
+			if (!text || ingestedContextTexts.has(text)) {
+				return;
+			}
+
+			ingestedContextTexts.add(text);
+			// Evict oldest entries (Set preserves insertion order) to bound memory.
+			while (ingestedContextTexts.size > MAX_INGESTED_CONTEXT_TEXTS) {
+				const oldest = ingestedContextTexts.values().next().value;
+				if (oldest === undefined) {
+					break;
+				}
+				ingestedContextTexts.delete(oldest);
+			}
+			ingestDocumentPage(text);
+		};
+
+		ingestContextText(context.fullPageContent);
+		ingestContextText(context.parentCommentContent);
+		for (const siblingCommentContent of context.siblingCommentsContents ?? []) {
+			ingestContextText(siblingCommentContent);
+		}
+
+		// Context arrived after word boundaries may already have fired. Re-send
+		// slow-lane context immediately so the next inference includes the thread.
+		if (currentView) {
+			slowLaneClient.updateContext(
+				buildSlowLaneText(currentView.state.doc.textContent, resolvedContext),
+			);
+		}
+	};
+
+	/**
+	 * Returns true when a fetch was actually started, false when it was skipped
+	 * (no getContext, a request already in flight, or throttled). Callers that
+	 * track a retry budget should only count attempts where this returns true.
+	 */
+	const refreshContext = ({
+		source,
+		allowThrottle = true,
+	}: {
+		allowThrottle?: boolean;
+		source: string;
+	}): boolean => {
+		if (!options?.getContext || contextRequestInFlight) {
+			return false;
+		}
+
+		const now = Date.now();
+		if (allowThrottle && now - lastContextRefreshAt < CONTEXT_REFRESH_THROTTLE_MS) {
+			return false;
+		}
+
+		contextRequestInFlight = true;
+		lastContextRefreshAt = now;
+
+		options
+			.getContext()
+			.then((context) => {
+				// Bail if the plugin was destroyed while the fetch was in flight —
+				// applyContext mutates global text-predictor state we must not touch
+				// after teardown.
+				if (destroyed) {
+					return;
+				}
+				logContextResolved(source, context);
+				if (!context) {
+					return;
+				}
+
+				applyContext(context);
+			})
+			.catch((error) => {
+				logException(error as Error, {
+					location: 'editor-plugin-autocomplete/getContext',
+				});
+			})
+			.finally(() => {
+				contextRequestInFlight = false;
+			});
+
+		return true;
+	};
+
 	/**
 	 * Schedule a prediction after a short debounce.
 	 * Tier 1 predictions are synchronous (<0.1ms) but we still debounce
@@ -490,29 +635,10 @@ export const createAutocompletePlugin = (
 					});
 					if (!hasIngestedPage) {
 						hasIngestedPage = true;
-						if (options?.getContext) {
-							options
-								.getContext()
-								.then((context) => {
-									if (!context) {
-										return;
-									}
-									resolvedContext = context;
-
-									if (context.fullPageContent) {
-										ingestDocumentPage(context.fullPageContent);
-									}
-									if (context.parentCommentContent) {
-										ingestDocumentPage(context.parentCommentContent);
-									}
-									context.siblingCommentsContents?.forEach(ingestDocumentPage);
-								})
-								.catch((error) => {
-									logException(error as Error, {
-										location: 'editor-plugin-autocomplete/getContext',
-									});
-								});
-						}
+						refreshContext({
+							source: 'focus',
+							allowThrottle: false,
+						});
 					}
 					return false;
 				},
@@ -521,6 +647,7 @@ export const createAutocompletePlugin = (
 
 		view: () => ({
 			update: (view: EditorView, prevState: EditorState) => {
+				currentView = view;
 				if (!prevState.doc.eq(view.state.doc)) {
 					if (justAccepted) {
 						justAccepted = false;
@@ -542,18 +669,36 @@ export const createAutocompletePlugin = (
 						slowLaneClient.updateContext(
 							buildSlowLaneText(view.state.doc.textContent, resolvedContext),
 						);
+
+						// Context may not have resolved on first focus (e.g. comment
+						// thread still loading). Retry on word boundaries until we have
+						// the parent comment, throttled so we don't refetch constantly
+						// and capped so non-comment editors stop retrying entirely.
+						if (
+							!resolvedContext?.parentCommentContent &&
+							wordBoundaryRefreshAttempts < MAX_CONTEXT_REFRESH_ATTEMPTS
+						) {
+							// Only count the attempt when a fetch actually started, so an
+							// in-flight or throttled no-op doesn't burn the retry budget.
+							if (refreshContext({ source: 'word-boundary' })) {
+								wordBoundaryRefreshAttempts++;
+							}
+						}
 					}
 
 					schedulePrediction(view);
 				}
 			},
 			destroy: () => {
+				destroyed = true;
+				currentView = null;
 				if (debounceTimer) {
 					clearTimeout(debounceTimer);
 				}
 				if (hasDestroy(slowLaneClient)) {
 					slowLaneClient.destroy();
 				}
+				ingestedContextTexts.clear();
 				setDefaultSlowLaneClient(null);
 			},
 		}),

package/src/pm-plugins/local-slow-lane-client.ts

@@ -41,6 +41,7 @@ import { isAutocompleteDebugEnabled } from './debug-mode';
 import { isWordBoundary } from './slow-lane-client';
 
 type WebLlmModelRecord = NonNullable<AppConfig['model_list']>[number];
+type EmbeddingApiResponse = { data?: Array<{ embedding?: unknown }> };
 
 // ─── Types ───────────────────────────────────────────────────────────────────
 
@@ -158,12 +159,42 @@ export const BE_PARITY = {
 	MAX_CONTEXT_WORDS: 100,
 } as const;
 
+const splitOnWhitespace = (text: string): string[] => {
+	const trimmed = text.trim();
+	if (trimmed === '') {
+		return [];
+	}
+
+	const words: string[] = [];
+	let wordStart = -1;
+
+	for (let i = 0; i < trimmed.length; i++) {
+		if (trimmed[i].trim() === '') {
+			if (wordStart !== -1) {
+				words.push(trimmed.slice(wordStart, i));
+				wordStart = -1;
+			}
+			continue;
+		}
+
+		if (wordStart === -1) {
+			wordStart = i;
+		}
+	}
+
+	if (wordStart !== -1) {
+		words.push(trimmed.slice(wordStart));
+	}
+
+	return words;
+};
+
 /**
  * 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: string, n: number): string => {
-	const words = text.trim().split(/\s+/u);
+	const words = splitOnWhitespace(text);
 	return words.length <= n ? text : words.slice(-n).join(' ');
 };
 
@@ -564,6 +595,14 @@ export const createLocalSlowLaneClient = (
 	let lastRequestedText = '';
 	let requestCounter = 0;
 	let latestRequestId = -1;
+	let inferenceInFlight = false;
+	let activeInferenceText: string | null = null;
+	// The requestId of the inference currently in flight. Tracked so the
+	// in-flight dedup path can restore `latestRequestId` to it — otherwise an
+	// intermediate keystroke that bumped `latestRequestId` would cause the
+	// in-flight (still-current) result to be discarded as stale.
+	let activeInferenceRequestId = -1;
+	let pendingInference: { requestId: number; text: string } | null = null;
 	let ready = false;
 	let destroyed = false;
 	let initFailed = false;
@@ -755,17 +794,25 @@ export const createLocalSlowLaneClient = (
 		const lmText = truncateToLastNWords(text, BE_PARITY.MAX_CONTEXT_TOKENS);
 		const semanticText = truncateToLastNWords(text, BE_PARITY.MAX_CONTEXT_WORDS);
 		const arcticInput = wrapForArctic(semanticText);
+		const captureCompletionTime = <T,>(
+			promise: Promise<T>,
+			onResolved: (resolvedAt: number) => void,
+		): Promise<T> =>
+			promise.then((value: T) => {
+				onResolved(performance.now());
+				return value;
+			});
 
 		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}"`,
+				`%c[LocalSlowLane] %c🔢 Arctic input (${arcticInput.length} chars, ${splitOnWhitespace(semanticText).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}"`,
+				`%c[LocalSlowLane] %c🧠 LM input (${lmText.length} chars, ${splitOnWhitespace(lmText).length} words): "${lmText.length > 100 ? `${lmText.slice(0, 100)}…` : lmText}"`,
 				'color: #9c27b0; font-weight: bold;',
 				'color: #2196f3;',
 			);
@@ -777,27 +824,29 @@ export const createLocalSlowLaneClient = (
 			let tEmbDone = 0;
 
 			const [, embeddingResponse] = await Promise.all([
-				engine.completions
-					.create({
+				captureCompletionTime(
+					engine.completions.create({
 						model: modelId,
 						prompt: lmText,
 						max_tokens: 1,
 						temperature: 0,
 						logprobs: false,
 					})
-					.then((r) => {
-						tLmDone = performance.now();
-						return r;
-					}),
-				engine.embeddings
-					.create({
+					,
+					(resolvedAt) => {
+						tLmDone = resolvedAt;
+					},
+				),
+				captureCompletionTime(
+					engine.embeddings.create({
 						model: LOCAL_MLC_EMBEDDING_MODEL_ID,
 						input: arcticInput,
 					})
-					.then((r) => {
-						tEmbDone = performance.now();
-						return r;
-					}),
+					,
+					(resolvedAt) => {
+						tEmbDone = resolvedAt;
+					},
+				),
 			]);
 
 			if (isAutocompleteDebugEnabled()) {
@@ -827,7 +876,7 @@ export const createLocalSlowLaneClient = (
 			// 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.data?.[0]?.embedding;
+			const embedding = (embeddingResponse as EmbeddingApiResponse).data?.[0]?.embedding;
 			storedContextVector =
 				Array.isArray(embedding) && embedding.length > 0
 					? new Float32Array(embedding as number[])
@@ -879,8 +928,8 @@ export const createLocalSlowLaneClient = (
 				hasLmLogits: storedLmLogits !== null,
 			});
 		} catch (err) {
-			// Discard errors for stale requests
-			if (requestId < latestRequestId) {
+			// Discard errors for stale requests or after teardown
+			if (requestId < latestRequestId || destroyed) {
 				return;
 			}
 
@@ -902,11 +951,55 @@ export const createLocalSlowLaneClient = (
 
 	// ── Context update (debounced) ─────────────────────────────────────────
 
+	const startInference = (text: string, requestId: number): void => {
+		// Self-contained guard: never start a new inference cycle after teardown,
+		// regardless of caller discipline.
+		if (destroyed) {
+			return;
+		}
+
+		inferenceInFlight = true;
+		activeInferenceText = text;
+		activeInferenceRequestId = requestId;
+
+		void ensureEngineInitialized()
+			.then(() => runInference(text, requestId))
+			.catch(() => {})
+			.finally(() => {
+				inferenceInFlight = false;
+				activeInferenceText = null;
+				activeInferenceRequestId = -1;
+
+				const next = pendingInference;
+				pendingInference = null;
+				if (next && !destroyed) {
+					startInference(next.text, next.requestId);
+				}
+			});
+	};
+
 	const doUpdateContext = (text: string): void => {
 		if (destroyed || !text || text.trim().length === 0) {
 			return;
 		}
 
+		if (inferenceInFlight && text === activeInferenceText) {
+			// The latest desired text already matches the in-flight inference, so
+			// re-running it would be wasted work. But an intermediate keystroke may
+			// have bumped `latestRequestId` past the in-flight request (and then been
+			// coalesced away), which would cause runInference to discard the
+			// still-current result as stale. Pin `latestRequestId` back to the active
+			// request so its result is accepted, and drop any now-superseded pending
+			// request.
+			latestRequestId = activeInferenceRequestId;
+			pendingInference = null;
+			return;
+		}
+
+		if (inferenceInFlight && pendingInference?.text === text) {
+			return;
+		}
+
 		const requestId = ++requestCounter;
 		latestRequestId = requestId;
 
@@ -926,15 +1019,26 @@ export const createLocalSlowLaneClient = (
 			console.groupEnd();
 		}
 
-		void ensureEngineInitialized()
-			.then(() => runInference(text, requestId))
-			.catch(() => {});
+		if (inferenceInFlight) {
+			pendingInference = { text, requestId };
+			return;
+		}
+
+		startInference(text, requestId);
 	};
 
 	const updateContextDebounced = (text: string): void => {
 		if (debounceTimer) {
 			clearTimeout(debounceTimer);
 		}
+		if (inferenceInFlight) {
+			pendingInference = null;
+			if (text === activeInferenceText) {
+				latestRequestId = activeInferenceRequestId;
+				lastRequestedText = text;
+				return;
+			}
+		}
 		lastRequestedText = text;
 		debounceTimer = setTimeout(() => {
 			debounceTimer = null;
@@ -967,6 +1071,10 @@ export const createLocalSlowLaneClient = (
 				unloadEngine(engineToUnload);
 			}
 			engineInitPromise = null;
+			inferenceInFlight = false;
+			activeInferenceText = null;
+			activeInferenceRequestId = -1;
+			pendingInference = null;
 			storedContextVector = null;
 			storedLmLogits = null;
 		},