From c192d6d7295a5f9b31ad98f0318f0cf6f8c5a9c2 Mon Sep 17 00:00:00 2001
From: Richard R <me@richardr.dev>
Date: Tue, 9 Jun 2026 15:35:39 -0600
Subject: [PATCH] fix(epub): improve word-region cache invalidation for
 highlighting

Update the EPUB highlighting cache key to include aligned word texts, ensuring
that changes in alignment or word content invalidate stale cached spans even if
word counts remain unchanged. Clear the word-region cache when rendered text
maps are replaced or reset, preventing incorrect highlights after remapping.

Also add locale tracking to HTML sentence highlighting state to support
language-aware word segmentation during alignment.
---
 src/hooks/epub/useEPUBHighlighting.ts | 11 ++++++++++-
 src/lib/client/html/highlight.ts      |  6 +++++-
 2 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/src/hooks/epub/useEPUBHighlighting.ts b/src/hooks/epub/useEPUBHighlighting.ts
index d186670..a9a7dc8 100644
--- a/src/hooks/epub/useEPUBHighlighting.ts
+++ b/src/hooks/epub/useEPUBHighlighting.ts
@@ -44,7 +44,9 @@ export function useEPUBHighlighting({
   renderedTextMapsRef,
 }: UseEpubHighlightingParams): UseEpubHighlightingResult {
   // Cache the per-segment word→region map so we don't re-align on every whisper
-  // tick. Keyed by the segment + resolved region + alignment.
+  // tick. Keyed by the segment + resolved region + the aligned word texts, so a
+  // corrected/rebuilt alignment (even with an identical word count) misses the
+  // cache instead of reusing stale spans.
   const wordRangeCacheRef = useRef<{ key: string; spans: Array<AlignmentCharSpan | null> } | null>(null);
 
   const clearWordHighlights = useCallback(() => {
@@ -122,6 +124,9 @@ export function useEPUBHighlighting({
       resolved.startOffset,
       resolved.endOffset,
       words.length,
+      // Word texts (not timings) drive the span mapping, so include them: a
+      // re-aligned segment with the same count still invalidates the cache.
+      words.map((word) => word.text).join('
'),
     ].join('::');
     if (wordRangeCacheRef.current?.key !== cacheKey) {
       wordRangeCacheRef.current = {
@@ -166,10 +171,14 @@ export function useEPUBHighlighting({
 
   const setRenderedTextMaps = useCallback((maps: EpubRenderedTextMap[]) => {
     renderedTextMapsRef.current = maps;
+    // Remapped content can change a region's text under an unchanged cache key,
+    // so drop the word-span cache whenever the text maps are replaced.
+    wordRangeCacheRef.current = null;
   }, [renderedTextMapsRef]);
 
   const resetHighlightState = useCallback(() => {
     renderedTextMapsRef.current = [];
+    wordRangeCacheRef.current = null;
     clearHighlights();
   }, [clearHighlights, renderedTextMapsRef]);
 
diff --git a/src/lib/client/html/highlight.ts b/src/lib/client/html/highlight.ts
index 23779df..668e4dc 100644
--- a/src/lib/client/html/highlight.ts
+++ b/src/lib/client/html/highlight.ts
@@ -66,6 +66,9 @@ interface SentenceState {
   // Normalized text of the wrap (chars joined), tokenized to align each spoken
   // word against the rendered words.
   text: string;
+  // Locale captured when the sentence was highlighted, reused for locale-aware
+  // word segmentation when mapping the alignment (matters for CJK/Thai, etc.).
+  language?: string;
   // For an alignment we've already seen: each word's [start, end) char span
   // within `text`/`chars` (null entries = words that aligned to no token).
   alignment: TTSSentenceAlignment | null;
@@ -318,6 +321,7 @@ export function highlightHtmlSentence(
     sentence,
     chars,
     text,
+    language,
     alignment: null,
     wordRanges: null,
   };
@@ -349,7 +353,7 @@ export function highlightHtmlWord(
   // rendered words and won't jump to a later/duplicate word.
   if (sentenceState.alignment !== alignment || !sentenceState.wordRanges) {
     sentenceState.alignment = alignment;
-    sentenceState.wordRanges = locateAlignmentWordSpans(words, sentenceState.text);
+    sentenceState.wordRanges = locateAlignmentWordSpans(words, sentenceState.text, sentenceState.language);
   }
 
   const range = sentenceState.wordRanges[wordIndex];