Add private personalized suggestion ranking

scripts/run-unit-tests.ts

@@ -4,8 +4,14 @@ const POPUP_TEST = "tests/popup.dashboard.retry.test.ts";
 const SUGGESTION_MANAGER_TEST = "tests/SuggestionManager.test.ts";
 
 const UTILS_TEST = "tests/utils.test.ts";
-
-const ISOLATED_TESTS = new Set([POPUP_TEST, SUGGESTION_MANAGER_TEST, UTILS_TEST]);
+const PERSONALIZATION_SERVICE_TEST = "tests/PersonalizationService.test.ts";
+
+const ISOLATED_TESTS = new Set([
+  POPUP_TEST,
+  SUGGESTION_MANAGER_TEST,
+  UTILS_TEST,
+  PERSONALIZATION_SERVICE_TEST,
+]);
 
 function sortedUnique(entries: string[]): string[] {
   return [...new Set(entries)].sort((left, right) => left.localeCompare(right));

src/adapters/chrome/background/BackgroundServiceWorker.ts

@@ -32,6 +32,10 @@ import {
 } from "./config/runtimeSettings";
 import { ConfigAssembler } from "./config/ConfigAssembler";
 import { ObservabilityService } from "./ObservabilityService";
+import { ChromeStorageBackend } from "@core/application/storage/ChromeStorageBackend";
+import { PersonalizationRepository } from "@core/application/personalization/PersonalizationRepository";
+import { PersonalizationService } from "@core/application/personalization/PersonalizationService";
+import type { PersonalizationEvent } from "@core/domain/personalization/types";
 
 declare const __FT_DEV_BUILD__: boolean | undefined;
 
@@ -49,6 +53,7 @@ export class BackgroundServiceWorker {
   productivityStatsManager!: ProductivityStatsManager;
   observabilityService!: ObservabilityService;
   configAssembler!: ConfigAssembler;
+  personalizationService!: PersonalizationService;
   language!: string;
   private runtimeConfigReady = false;
   private runtimeConfigLoadPromise: Promise<void> | null = null;
@@ -60,8 +65,21 @@ export class BackgroundServiceWorker {
     }
     this.settingsManager = new SettingsManager();
     this.coreSettingsRepository = new CoreSettingsRepository(this.settingsManager);
+    this.personalizationService = new PersonalizationService({
+      repository: new PersonalizationRepository(new ChromeStorageBackend(true)),
+      isEnabled: () => this.coreSettingsRepository.getPersonalizationEnabled(),
+      isTextExpansionTrigger: async (triggerText) => {
+        const normalizedTrigger = triggerText.trim().toLocaleLowerCase();
+        const expansions = await this.coreSettingsRepository.getTextExpansions();
+        return expansions.some(
+          ([shortcut]) => shortcut.trim().toLocaleLowerCase() === normalizedTrigger,
+        );
+      },
+    });
     this.languageDetector = new LanguageDetector(this.settingsManager);
-    this.predictionManager = new PredictionManager();
+    this.predictionManager = new PredictionManager({
+      getPersonalizationSnapshot: () => this.personalizationService.getRankingSnapshot(),
+    });
     this.tabMessenger = new TabMessenger();
     this.productivityStatsManager = new ProductivityStatsManager(this.settingsManager);
     this.observabilityService = new ObservabilityService({
@@ -216,7 +234,10 @@ export class BackgroundServiceWorker {
   async updatePresageConfig(): Promise<void> {
     await sanitizeSiteProfilesSetting(this.settingsManager);
     await sanitizeAutoLanguagePriorsSetting(this.settingsManager);
-    await this.predictionManager.initialize();
+    await Promise.all([
+      this.personalizationService.initialize(),
+      this.predictionManager.initialize(),
+    ]);
     const runtimeConfig = await this.configAssembler.assemblePredictionRuntimeConfig();
     this.language = runtimeConfig.language;
     this.observabilityService.setConfig(runtimeConfig.observabilityConfig);
@@ -304,7 +325,10 @@ export class BackgroundServiceWorker {
         await migrateSettingsV5(this.settingsManager);
         await migrateSettingsV6(this.settingsManager);
         await migrateSettingsV7(this.settingsManager);
-        await this.predictionManager.initialize();
+        await Promise.all([
+          this.personalizationService.initialize(),
+          this.predictionManager.initialize(),
+        ]);
         await this.updatePresageConfig();
       } catch (error) {
         logError("lastVersion handler", error);
@@ -313,6 +337,14 @@ export class BackgroundServiceWorker {
     await this.initializationPromise;
   }
 
+  async handlePersonalizationEvent(event: PersonalizationEvent): Promise<boolean> {
+    return this.personalizationService.handleEvent(event);
+  }
+
+  async clearPersonalization(): Promise<void> {
+    await this.personalizationService.clear();
+  }
+
   private async ensureRuntimeConfigReady(): Promise<void> {
     if (this.runtimeConfigReady) {
       return;

src/adapters/chrome/background/PredictionManager.ts

@@ -15,6 +15,11 @@ import { WebLLMPredictor } from "./WebLLMPredictor";
 import { createLogger } from "@core/application/logging/Logger";
 import { DEFAULT_AI_PREDICTION_TIMEOUT_MS } from "@core/domain/constants";
 import { PredictorError, getErrorMessage } from "@core/domain/error";
+import type { PersonalizationRankingSnapshot } from "@core/domain/personalization/types";
+
+interface PredictionManagerOptions {
+  getPersonalizationSnapshot?: () => PersonalizationRankingSnapshot;
+}
 
 export interface PredictionDebugRequestMeta {
   traceId?: string;
@@ -97,9 +102,11 @@ export class PredictionManager {
   private debugTraces: PredictorDebugTrace[] = [];
   private debugTraceById: Map<string, PredictorDebugTrace> = new Map();
   private currentConfig: PredictionConfig | null = null;
+  private readonly getPersonalizationSnapshot: () => PersonalizationRankingSnapshot;
 
-  constructor() {
+  constructor(options: PredictionManagerOptions = {}) {
     this.libPresageMod = libPresageMod as () => Promise<PresageModule>;
+    this.getPersonalizationSnapshot = options.getPersonalizationSnapshot ?? (() => ({}));
     void this.initialize();
   }
 
@@ -113,7 +120,9 @@ export class PredictionManager {
   private async _doInitializePresage(): Promise<void> {
     try {
       const Module = await this.libPresageMod();
-      this.presageHandler = new PresageHandler(Module);
+      this.presageHandler = new PresageHandler(Module, {
+        getPersonalizationSnapshot: this.getPersonalizationSnapshot,
+      });
       this.predictionOrchestrator = new PredictionOrchestrator(
         this.presageHandler,
         this.getWebLLMPredictor(),

src/adapters/chrome/background/PresageHandler.ts

@@ -13,6 +13,8 @@ import { PresageEngine } from "./PresageEngine";
 import { MAX_NUM_SUGGESTIONS } from "@core/domain/constants";
 import type { PredictionResult } from "./PredictionTypes";
 import { SPACING_RULES, Spacing } from "@core/domain/spacingRules";
+import { rankPersonalizedCandidates } from "@core/domain/personalization/PersonalizationRanker";
+import type { PersonalizationRankingSnapshot } from "@core/domain/personalization/types";
 const SUGGESTION_COUNT = 5;
 const MIN_WORD_LENGTH_TO_PREDICT = 1;
 const logger = createLogger("PresageHandler");
@@ -25,12 +27,18 @@ export interface PresageConfig {
   autoCapitalize: boolean;
   textExpansions: Array<[string, object]>;
   prefixOnlyMode: boolean;
+  personalizationEnabled?: boolean;
 
   timeFormat?: string;
   dateFormat?: string;
   userDictionaryList?: string[];
 }
 
+interface PresageHandlerOptions {
+  getPersonalizationSnapshot?: () => PersonalizationRankingSnapshot;
+  now?: () => number;
+}
+
 export interface PresagePredictionContext {
   text: string;
   nextChar: string;
@@ -61,9 +69,13 @@ export class PresageHandler {
   private dateFormat?: string;
   private engineNumSuggestions: number;
   private textExpansionsSignature = "";
+  private textExpansionShortcuts = new Set<string>();
   private userDictionarySignature = "";
+  private personalizationEnabled = false;
+  private readonly getPersonalizationSnapshot: () => PersonalizationRankingSnapshot;
+  private readonly now: () => number;
 
-  constructor(Module: PresageModule) {
+  constructor(Module: PresageModule, options: PresageHandlerOptions = {}) {
     const engineConfig: PresageEngineConfig = {
       numSuggestions: SUGGESTION_COUNT,
       prefixOnlyMode: false,
@@ -77,6 +89,8 @@ export class PresageHandler {
     this.autoCapitalize = true;
     this.prefixOnlyMode = false;
     this.userDictionaryList = [];
+    this.getPersonalizationSnapshot = options.getPersonalizationSnapshot ?? (() => ({}));
+    this.now = options.now ?? Date.now;
 
     this.predictionInputProcessor = new PredictionInputProcessor(
       this.minWordLengthToPredict,
@@ -116,6 +130,10 @@ export class PresageHandler {
     this.insertSpaceAfterAutocomplete = config.insertSpaceAfterAutocomplete;
     this.autoCapitalize = config.autoCapitalize;
     this.prefixOnlyMode = config.prefixOnlyMode;
+    this.personalizationEnabled = config.personalizationEnabled ?? false;
+    this.textExpansionShortcuts = new Set(
+      (config.textExpansions ?? []).map(([shortcut]) => shortcut.trim().toLocaleLowerCase()),
+    );
 
     this.timeFormat = config.timeFormat;
     this.dateFormat = config.dateFormat;
@@ -235,7 +253,26 @@ export class PresageHandler {
     ) {
       return [];
     }
-    return this.doPredictionHandler(context.predictionInput, context.lang, context.tabId);
+    const predictions = await this.doPredictionHandler(
+      context.predictionInput,
+      context.lang,
+      context.tabId,
+    );
+    if (!this.personalizationEnabled || this.isTextExpansionRequest(context.predictionInput)) {
+      return predictions;
+    }
+
+    const inputLower = context.predictionInput.trim().toLocaleLowerCase();
+    const pinnedCandidates = new Set(
+      predictions.filter((candidate) => candidate.toLocaleLowerCase() === inputLower),
+    );
+    return rankPersonalizedCandidates({
+      candidates: predictions,
+      language: context.lang,
+      snapshot: this.getPersonalizationSnapshot(),
+      nowMs: this.now(),
+      pinnedCandidates,
+    });
   }
 
   finalizePrediction(
@@ -334,4 +371,9 @@ export class PresageHandler {
       presageEngine.reinitialize();
     }
   }
+
+  private isTextExpansionRequest(predictionInput: string): boolean {
+    const finalToken = predictionInput.trim().split(/\s+/u).at(-1)?.toLocaleLowerCase();
+    return finalToken ? this.textExpansionShortcuts.has(finalToken) : false;
+  }
 }

src/adapters/chrome/background/config/ConfigAssembler.ts

@@ -109,6 +109,7 @@ export class ConfigAssembler {
       observability,
       prefixOnlyMode,
       inlineSuggestion,
+      personalizationEnabled,
     ] = await Promise.all([
       this.coreSettingsRepository.getNumSuggestions(),
       this.coreSettingsRepository.getMinWordLengthToPredict(),
@@ -123,6 +124,7 @@ export class ConfigAssembler {
       this.getObservabilityConfig(),
       this.coreSettingsRepository.getPrefixOnlyMode(),
       this.coreSettingsRepository.getInlineSuggestion(),
+      this.coreSettingsRepository.getPersonalizationEnabled(),
     ]);
     const normalizedGrammarRules = normalizeGrammarRuleSelection(enabledGrammarRules);
     const autoCapitalize = normalizedGrammarRules.includes("capitalizeSentenceStart");
@@ -139,6 +141,7 @@ export class ConfigAssembler {
         autoCapitalize,
         textExpansions,
         prefixOnlyMode: prefixOnlyMode || inlineSuggestion,
+        personalizationEnabled,
 
         timeFormat,
         dateFormat,

src/adapters/chrome/background/router/MessageRouter.ts

@@ -7,6 +7,7 @@ import {
   CMD_CONTENT_SCRIPT_REPORT_OBSERVABILITY_MODULES,
   CMD_CONTENT_SCRIPT_REPORT_RUNTIME_STATUS,
   CMD_CONTENT_SCRIPT_USAGE_EVENT,
+  CMD_CONTENT_SCRIPT_PERSONALIZATION_EVENT,
   CMD_GET_AUTO_LANGUAGE_STATUS,
   CMD_OPTIONS_CLEAR_OBSERVABILITY_EVENTS,
   CMD_OPTIONS_CLEAR_PREDICTOR_DEBUG_TRACE,
@@ -16,6 +17,7 @@ import {
   CMD_OPTIONS_REPORT_OBSERVABILITY_MODULES,
   CMD_OPTIONS_PAGE_CONFIG_CHANGE,
   CMD_OPTIONS_RESET_PRODUCTIVITY_STATS,
+  CMD_OPTIONS_CLEAR_PERSONALIZATION,
   CMD_POPUP_ACK_DONATION_MILESTONE,
   CMD_POPUP_ACK_WEEKLY_RECAP,
   CMD_POPUP_GET_PRODUCTIVITY_STATS,
@@ -53,6 +55,7 @@ const ROUTED_MESSAGE_COMMANDS = [
   CMD_OPTIONS_PAGE_CONFIG_CHANGE,
   CMD_CONTENT_SCRIPT_GET_CONFIG,
   CMD_CONTENT_SCRIPT_USAGE_EVENT,
+  CMD_CONTENT_SCRIPT_PERSONALIZATION_EVENT,
   CMD_CONTENT_SCRIPT_REPORT_RUNTIME_STATUS,
   CMD_CONTENT_SCRIPT_REPORT_OBSERVABILITY_EVENT,
   CMD_CONTENT_SCRIPT_REPORT_OBSERVABILITY_MODULES,
@@ -61,6 +64,7 @@ const ROUTED_MESSAGE_COMMANDS = [
   CMD_POPUP_ACK_WEEKLY_RECAP,
   CMD_POPUP_ACK_DONATION_MILESTONE,
   CMD_OPTIONS_RESET_PRODUCTIVITY_STATS,
+  CMD_OPTIONS_CLEAR_PERSONALIZATION,
   CMD_OPTIONS_GET_PREDICTOR_DEBUG_SNAPSHOT,
   CMD_OPTIONS_CLEAR_PREDICTOR_DEBUG_TRACE,
   CMD_OPTIONS_GET_OBSERVABILITY_SNAPSHOT,
@@ -122,6 +126,8 @@ const MESSAGE_ERROR_LABELS: Record<RoutedMessageCommand, string> = {
   [CMD_OPTIONS_PAGE_CONFIG_CHANGE]: "handleOptionsPageConfigChange",
   [CMD_CONTENT_SCRIPT_GET_CONFIG]: "MessageRouter.handleContentScriptGetConfig",
   [CMD_CONTENT_SCRIPT_USAGE_EVENT]: "MessageRouter.handleContentScriptUsageEvent",
+  [CMD_CONTENT_SCRIPT_PERSONALIZATION_EVENT]:
+    "MessageRouter.handleContentScriptPersonalizationEvent",
   [CMD_CONTENT_SCRIPT_REPORT_RUNTIME_STATUS]: "MessageRouter.handleContentScriptRuntimeStatus",
   [CMD_CONTENT_SCRIPT_REPORT_OBSERVABILITY_EVENT]:
     "MessageRouter.handleContentScriptReportObservabilityEvent",
@@ -132,6 +138,7 @@ const MESSAGE_ERROR_LABELS: Record<RoutedMessageCommand, string> = {
   [CMD_POPUP_ACK_WEEKLY_RECAP]: "MessageRouter.handlePopupAckWeeklyRecap",
   [CMD_POPUP_ACK_DONATION_MILESTONE]: "MessageRouter.handlePopupAckDonationMilestone",
   [CMD_OPTIONS_RESET_PRODUCTIVITY_STATS]: "MessageRouter.handleOptionsResetProductivityStats",
+  [CMD_OPTIONS_CLEAR_PERSONALIZATION]: "MessageRouter.handleOptionsClearPersonalization",
   [CMD_OPTIONS_GET_PREDICTOR_DEBUG_SNAPSHOT]:
     "MessageRouter.handleOptionsGetPredictorDebugSnapshot",
   [CMD_OPTIONS_CLEAR_PREDICTOR_DEBUG_TRACE]: "MessageRouter.handleOptionsClearPredictorDebugTrace",
@@ -180,6 +187,10 @@ export class MessageRouter {
     register(CMD_OPTIONS_PAGE_CONFIG_CHANGE, this.handleOptionsPageConfigChange.bind(this));
     register(CMD_CONTENT_SCRIPT_GET_CONFIG, this.handleContentScriptGetConfig.bind(this));
     register(CMD_CONTENT_SCRIPT_USAGE_EVENT, this.handleContentScriptUsageEvent.bind(this));
+    register(
+      CMD_CONTENT_SCRIPT_PERSONALIZATION_EVENT,
+      this.handleContentScriptPersonalizationEvent.bind(this),
+    );
     register(
       CMD_CONTENT_SCRIPT_REPORT_RUNTIME_STATUS,
       this.handleContentScriptRuntimeStatus.bind(this),
@@ -200,6 +211,7 @@ export class MessageRouter {
       CMD_OPTIONS_RESET_PRODUCTIVITY_STATS,
       this.handleOptionsResetProductivityStats.bind(this),
     );
+    register(CMD_OPTIONS_CLEAR_PERSONALIZATION, this.handleOptionsClearPersonalization.bind(this));
     register(
       CMD_OPTIONS_GET_PREDICTOR_DEBUG_SNAPSHOT,
       this.handleOptionsGetPredictorDebugSnapshot.bind(this),
@@ -441,6 +453,14 @@ export class MessageRouter {
     this.respondOk(sendResponse);
   }
 
+  private async handleContentScriptPersonalizationEvent(
+    payload: CommandPayload<typeof CMD_CONTENT_SCRIPT_PERSONALIZATION_EVENT>,
+  ): Promise<void> {
+    const { request, sendResponse, worker } = payload;
+    await worker.handlePersonalizationEvent(request.context);
+    this.respondOk(sendResponse);
+  }
+
   private handleContentScriptRuntimeStatus(
     payload: CommandPayload<typeof CMD_CONTENT_SCRIPT_REPORT_RUNTIME_STATUS>,
   ): void {
@@ -555,6 +575,14 @@ export class MessageRouter {
     this.respondOk(sendResponse);
   }
 
+  private async handleOptionsClearPersonalization(
+    payload: CommandPayload<typeof CMD_OPTIONS_CLEAR_PERSONALIZATION>,
+  ): Promise<void> {
+    const { sendResponse, worker } = payload;
+    await worker.clearPersonalization();
+    this.respondOk(sendResponse);
+  }
+
   private async handleOptionsGetPredictorDebugSnapshot(
     payload: CommandPayload<typeof CMD_OPTIONS_GET_PREDICTOR_DEBUG_SNAPSHOT>,
   ): Promise<void> {

src/adapters/chrome/content-script/suggestions/SuggestionEntrySession.ts

@@ -60,6 +60,9 @@ export class SuggestionEntrySession {
   private readonly renderInline: () => void;
   private readonly recordSuggestionShown: SuggestionEntrySessionOptions["recordSuggestionShown"];
   private readonly recordSuggestionAccepted: SuggestionEntrySessionOptions["recordSuggestionAccepted"];
+  private readonly recordPersonalizationAccepted: NonNullable<
+    SuggestionEntrySessionOptions["recordPersonalizationAccepted"]
+  >;
   private readonly getLang: () => string;
   private readonly insertSpaceAfterAutocomplete: boolean;
   private readonly logRenderedSuggestionPopup: SuggestionEntrySessionOptions["logRenderedSuggestionPopup"];
@@ -84,6 +87,7 @@ export class SuggestionEntrySession {
     this.renderInline = options.renderInline;
     this.recordSuggestionShown = options.recordSuggestionShown;
     this.recordSuggestionAccepted = options.recordSuggestionAccepted;
+    this.recordPersonalizationAccepted = options.recordPersonalizationAccepted ?? (() => "");
     this.getLang = options.getLang;
     this.insertSpaceAfterAutocomplete = options.insertSpaceAfterAutocomplete;
     this.logRenderedSuggestionPopup = options.logRenderedSuggestionPopup;
@@ -1242,6 +1246,14 @@ export class SuggestionEntrySession {
       return false;
     }
     this.lastAcceptedSuggestion = suggestion;
+    const personalizationEventId = this.recordPersonalizationAccepted({
+      suggestion,
+      triggerText: accepted.triggerText,
+      language: this.getLang(),
+    });
+    if (personalizationEventId && this.entry.pendingExtensionEdit?.source === "suggestion") {
+      this.entry.pendingExtensionEdit.personalizationEventId = personalizationEventId;
+    }
     this.finishAcceptedSuggestion(
       accepted.triggerText,
       accepted.insertedText,