perf(subtitles): add concurrent chunk processing for AI segmentation

src/entrypoints/subtitles.content/segmentation-pipeline.ts

@@ -1,8 +1,13 @@
 import type { SubtitlesFragment } from "@/utils/subtitles/types"
 import { getLocalConfig } from "@/utils/config/storage"
-import { PROCESS_LOOK_AHEAD_MS } from "@/utils/constants/subtitles"
+import { MAX_CONCURRENT_SEGMENTS, PROCESS_LOOK_AHEAD_MS } from "@/utils/constants/subtitles"
 import { aiSegmentBlock } from "@/utils/subtitles/processor/ai-segmentation"
-import { optimizeSubtitles } from "@/utils/subtitles/processor/optimizer"
+import { optimizeSubtitles, rebalanceToTargetRange } from "@/utils/subtitles/processor/optimizer"
+
+interface ChunkResult {
+  fragments: SubtitlesFragment[]
+  chunk: SubtitlesFragment[]
+}
 
 export class SegmentationPipeline {
   // Segmented results, read by translation pipeline
@@ -16,15 +21,18 @@ export class SegmentationPipeline {
 
   private getVideoElement: () => HTMLVideoElement | null
   private getSourceLanguage: () => string
+  private onChunkProcessed: (() => void) | null
 
   constructor(options: {
     rawFragments: SubtitlesFragment[]
     getVideoElement: () => HTMLVideoElement | null
     getSourceLanguage: () => string
+    onChunkProcessed?: () => void
   }) {
     this.rawFragments = options.rawFragments
     this.getVideoElement = options.getVideoElement
     this.getSourceLanguage = options.getSourceLanguage
+    this.onChunkProcessed = options.onChunkProcessed ?? null
   }
 
   get isRunning(): boolean {
@@ -68,59 +76,102 @@ export class SegmentationPipeline {
 
     try {
       while (!this.stopped && this.hasUnprocessedChunks()) {
-        const didWork = await this.processNextChunk(video.currentTime * 1000)
-        if (!didWork)
+        const currentTimeMs = video.currentTime * 1000
+        const chunks = this.findNextChunks(currentTimeMs, MAX_CONCURRENT_SEGMENTS)
+        if (chunks.length === 0)
+          break
+
+        // Mark all fragments as in-progress
+        for (const chunk of chunks) {
+          chunk.forEach(f => this.segmentedRawStarts.add(f.start))
+        }
+
+        // Process chunks concurrently, then merge results synchronously
+        const results = await Promise.all(chunks.map(chunk => this.processChunk(chunk)))
+        if (this.stopped) {
+          // Roll back claimed fragments so they can be reprocessed on restart
+          for (const chunk of chunks) {
+            chunk.forEach(f => this.segmentedRawStarts.delete(f.start))
+          }
           break
+        }
+        for (const result of results) {
+          this.mergeFragments(result.fragments, result.chunk)
+        }
+        try {
+          this.onChunkProcessed?.()
+        }
+        catch {
+          // callback errors must not kill the segmentation loop
+        }
       }
     }
     finally {
       this.running = false
     }
   }
 
-  private async processNextChunk(currentTimeMs: number): Promise<boolean> {
-    const chunk = this.findNextChunk(currentTimeMs)
-    if (chunk.length === 0)
-      return false
-
-    chunk.forEach(f => this.segmentedRawStarts.add(f.start))
-
+  private async processChunk(chunk: SubtitlesFragment[]): Promise<ChunkResult> {
     try {
       const config = await getLocalConfig()
       if (config) {
         const segmented = await aiSegmentBlock(chunk, config)
-        const optimized = optimizeSubtitles(segmented, this.getSourceLanguage())
-        const chunkStart = chunk[0].start
-        const chunkEnd = chunk.at(-1)!.end
-        this.processedFragments = this.processedFragments.filter(
-          f => f.start < chunkStart || f.start > chunkEnd,
-        )
-        this.processedFragments.push(...optimized)
-        this.processedFragments.sort((a, b) => a.start - b.start)
+        const rebalanced = rebalanceToTargetRange(segmented, this.getSourceLanguage())
+        return { fragments: rebalanced, chunk }
       }
     }
-    catch {
+    catch (error) {
+      console.warn("[SegmentationPipeline] AI segmentation failed, falling back:", error)
       chunk.forEach(f => this.aiSegmentFailedRawStarts.add(f.start))
       const optimized = optimizeSubtitles(chunk, this.getSourceLanguage())
-      this.processedFragments.push(...optimized)
-      this.processedFragments.sort((a, b) => a.start - b.start)
+      return { fragments: optimized, chunk }
+    }
+
+    // Config unavailable — fall back to non-AI processing to avoid dropping chunks
+    const optimized = optimizeSubtitles(chunk, this.getSourceLanguage())
+    return { fragments: optimized, chunk }
+  }
+
+  private mergeFragments(newFragments: SubtitlesFragment[], chunk: SubtitlesFragment[]): void {
+    const rawStarts = new Set(chunk.map(f => f.start))
+    this.processedFragments = this.processedFragments.filter(
+      f => !rawStarts.has(f.start),
+    )
+    this.processedFragments.push(...newFragments)
+    this.processedFragments.sort((a, b) => a.start - b.start)
+  }
+
+  /**
+   * Find up to `maxChunks` non-overlapping chunks, prioritizing fragments
+   * closest to the current playback position.
+   */
+  private findNextChunks(currentTimeMs: number, maxChunks: number): SubtitlesFragment[][] {
+    const chunks: SubtitlesFragment[][] = []
+    const claimed = new Set<number>()
+
+    for (let i = 0; i < maxChunks; i++) {
+      const chunk = this.findNextChunk(currentTimeMs, claimed)
+      if (chunk.length === 0)
+        break
+      chunks.push(chunk)
+      chunk.forEach(f => claimed.add(f.start))
     }
 
-    return true
+    return chunks
   }
 
-  private findNextChunk(currentTimeMs: number): SubtitlesFragment[] {
+  private findNextChunk(currentTimeMs: number, claimed: Set<number>): SubtitlesFragment[] {
     const searchStart = Math.max(0, currentTimeMs - 10_000)
     const firstUnprocessed = this.rawFragments.find(
-      f => f.start >= searchStart && !this.segmentedRawStarts.has(f.start),
+      f => f.start >= searchStart && !this.segmentedRawStarts.has(f.start) && !claimed.has(f.start),
     )
     if (!firstUnprocessed)
       return []
 
     const windowEnd = firstUnprocessed.start + PROCESS_LOOK_AHEAD_MS
     return this.rawFragments.filter(
       f => f.start >= firstUnprocessed.start && f.start < windowEnd
-        && !this.segmentedRawStarts.has(f.start),
+        && !this.segmentedRawStarts.has(f.start) && !claimed.has(f.start),
     )
   }
 }

src/entrypoints/subtitles.content/translation-coordinator.ts

@@ -19,6 +19,7 @@ export class TranslationCoordinator {
   private translatedStarts = new Set<number>()
   private failedStarts = new Set<number>()
   private isTranslating = false
+  private stopped = false
   private lastEmittedState: SubtitlesState = "idle"
   private videoContext: SubtitlesVideoContext = { videoTitle: "", subtitlesTextContent: "" }
 
@@ -43,6 +44,8 @@ export class TranslationCoordinator {
       this.videoContext = videoContext
     }
 
+    this.stopped = false
+
     const video = this.getVideoElement()
     if (!video)
       return
@@ -59,6 +62,7 @@ export class TranslationCoordinator {
   }
 
   stop() {
+    this.stopped = true
     const video = this.getVideoElement()
     if (!video)
       return
@@ -81,6 +85,12 @@ export class TranslationCoordinator {
     this.failedStarts.clear()
   }
 
+  triggerTranslationTick() {
+    if (this.stopped)
+      return
+    this.handleTranslationTick()
+  }
+
   private handleTranslationTick = () => {
     const video = this.getVideoElement()
     if (!video)

src/entrypoints/subtitles.content/universal-adapter.ts

@@ -332,6 +332,7 @@ export class UniversalVideoAdapter {
         rawFragments: this.originalSubtitles,
         getVideoElement: () => this.subtitlesScheduler?.getVideoElement() ?? null,
         getSourceLanguage: () => this.subtitlesFetcher.getSourceLanguage(),
+        onChunkProcessed: () => this.translationCoordinator?.triggerTranslationTick(),
       })
     }
     else {

src/utils/constants/subtitles.ts

@@ -10,6 +10,9 @@ export const MAX_WORDS = 15
 export const MAX_CHARS_CJK = 30
 export const SENTENCE_END_PATTERN = /[,.。?？！!；;…؟۔\n]$/
 
+// AI segmentation concurrency
+export const MAX_CONCURRENT_SEGMENTS = 3
+
 // On-demand translation constants
 export const TRANSLATION_BATCH_SIZE = 5
 export const TRANSLATE_LOOK_AHEAD_MS = 30_000

src/utils/subtitles/processor/optimizer.ts

@@ -165,7 +165,7 @@ function shouldKeepBoundary(left: SubtitlesFragment, right: SubtitlesFragment):
   return isTimeout || startsWithSign
 }
 
-function rebalanceToTargetRange(
+export function rebalanceToTargetRange(
   fragments: SubtitlesFragment[],
   language: string,
 ): SubtitlesFragment[] {