Merge pull request #24468 from overleaf/mj-client-side-references

[web] Perform ARS on client-side GitOrigin-RevId: 19703c82758cae450fe52463ad9612d3a2383ba0
2026-11-14 14:13:43 +01:00 · 2025-09-03 08:05:41 +00:00
parent f6820ed794
commit ed0c4c447e
21 changed files with 3040 additions and 26 deletions
@@ -87,13 +87,13 @@ export const ReactContextRoot: FC<
                            <Providers.EditorProvider>
                              <Providers.FileTreeDataProvider>
                                <Providers.FileTreePathProvider>
-                                  <Providers.ReferencesProvider>
-                                    <Providers.UserFeaturesProvider>
-                                      <Providers.PermissionsProvider>
-                                        <Providers.RailProvider>
-                                          <Providers.LayoutProvider>
-                                            <Providers.ProjectSettingsProvider>
-                                              <Providers.EditorManagerProvider>
+                                  <Providers.UserFeaturesProvider>
+                                    <Providers.PermissionsProvider>
+                                      <Providers.RailProvider>
+                                        <Providers.LayoutProvider>
+                                          <Providers.ProjectSettingsProvider>
+                                            <Providers.EditorManagerProvider>
+                                              <Providers.ReferencesProvider>
                                                <Providers.LocalCompileProvider>
                                                  <Providers.DetachCompileProvider>
                                                    <Providers.ChatProvider>
@@ -113,13 +113,13 @@ export const ReactContextRoot: FC<
                                                    </Providers.ChatProvider>
                                                  </Providers.DetachCompileProvider>
                                                </Providers.LocalCompileProvider>
-                                              </Providers.EditorManagerProvider>
-                                            </Providers.ProjectSettingsProvider>
-                                          </Providers.LayoutProvider>
-                                        </Providers.RailProvider>
-                                      </Providers.PermissionsProvider>
-                                    </Providers.UserFeaturesProvider>
-                                  </Providers.ReferencesProvider>
+                                              </Providers.ReferencesProvider>
+                                            </Providers.EditorManagerProvider>
+                                          </Providers.ProjectSettingsProvider>
+                                        </Providers.LayoutProvider>
+                                      </Providers.RailProvider>
+                                    </Providers.PermissionsProvider>
+                                  </Providers.UserFeaturesProvider>
                                </Providers.FileTreePathProvider>
                              </Providers.FileTreeDataProvider>
                            </Providers.EditorProvider>
@@ -7,21 +7,31 @@ import {
  useCallback,
  useMemo,
  useState,
+  useRef,
 } from 'react'
 import { useIdeReactContext } from '@/features/ide-react/context/ide-react-context'
 import { useConnectionContext } from '@/features/ide-react/context/connection-context'
-import { postJSON } from '@/infrastructure/fetch-json'
 import { ShareJsDoc } from '@/features/ide-react/editor/share-js-doc'
 import { useFileTreeData } from '@/shared/context/file-tree-data-context'
 import { findDocEntityById } from '@/features/ide-react/util/find-doc-entity-by-id'
 import { IdeEvents } from '@/features/ide-react/create-ide-event-emitter'
-import { debugConsole } from '@/utils/debugging'
 import useEventListener from '@/shared/hooks/use-event-listener'
+import { useProjectContext } from '@/shared/context/project-context'
+import { useEditorManagerContext } from './editor-manager-context'
+import { signalWithTimeout } from '@/utils/abort-signal'
+import { postJSON } from '@/infrastructure/fetch-json'
+import { debugConsole } from '@/utils/debugging'
+import { useFeatureFlag } from '@/shared/context/split-test-context'
+import type { ReferenceIndexer } from '../references/reference-indexer'
+import { AdvancedReferenceSearchResult } from '@/features/ide-react/references/types'

 export const ReferencesContext = createContext<
  | {
      referenceKeys: Set<string>
      indexAllReferences: (shouldBroadcast: boolean) => Promise<void>
+      searchLocalReferences: (
+        query: string
+      ) => Promise<AdvancedReferenceSearchResult>
    }
  | undefined
 >(undefined)
@@ -32,14 +42,18 @@ export const ReferencesProvider: FC<React.PropsWithChildren> = ({
  const { fileTreeData } = useFileTreeData()
  const { eventEmitter, projectId } = useIdeReactContext()
  const { socket } = useConnectionContext()
+  const { projectSnapshot } = useProjectContext()
+  const { openDocs } = useEditorManagerContext()
+  const abortControllerRef = useRef<AbortController | null>(null)

  const [referenceKeys, setReferenceKeys] = useState(new Set<string>())
+  const clientSideReferences = useFeatureFlag('client-side-references')

  const [existingIndexHash, setExistingIndexHash] = useState<
    Record<string, { hash: string; timestamp: number }>
  >({})

-  const indexAllReferences = useCallback(
+  const indexAllReferencesServerside = useCallback(
    async (shouldBroadcast: boolean) => {
      return postJSON(`/project/${projectId}/references/indexAll`, {
        body: {
@@ -57,6 +71,54 @@ export const ReferencesProvider: FC<React.PropsWithChildren> = ({
    [projectId]
  )

+  const indexerRef = useRef<Promise<ReferenceIndexer> | null>(null)
+  if (clientSideReferences && indexerRef.current === null) {
+    indexerRef.current = import('../references/reference-indexer').then(
+      m => new m.ReferenceIndexer()
+    )
+  }
+
+  const indexAllReferencesLocally = useCallback(
+    async (shouldBroadcast: boolean) => {
+      abortControllerRef.current?.abort()
+
+      if (!indexerRef.current) {
+        return
+      }
+
+      abortControllerRef.current = new AbortController()
+      const signal = abortControllerRef.current.signal
+
+      await openDocs.awaitBufferedOps(signalWithTimeout(signal, 5000))
+      await projectSnapshot.refresh()
+
+      if (signal.aborted) {
+        return
+      }
+
+      const indexer = await indexerRef.current
+      const keys = await indexer.updateFromSnapshot(projectSnapshot, { signal })
+      if (signal.aborted) {
+        return
+      }
+      setReferenceKeys(keys)
+      if (shouldBroadcast) {
+        // Inform other clients about change in keys
+        await postJSON(`/project/${projectId}/references/indexAll`, {
+          body: { shouldBroadcast: true },
+        }).catch(error => {
+          // allow the request to fail
+          debugConsole.error(error)
+        })
+      }
+    },
+    [projectSnapshot, openDocs, projectId]
+  )
+
+  const indexAllReferences = clientSideReferences
+    ? indexAllReferencesLocally
+    : indexAllReferencesServerside
+
  const indexReferencesIfDocModified = useCallback(
    (doc: ShareJsDoc, shouldBroadcast: boolean) => {
      // avoid reindexing references if the bib file has not changed since the
@@ -115,9 +177,13 @@ export const ReferencesProvider: FC<React.PropsWithChildren> = ({
      // We only need to grab the references when the editor first loads,
      // not on every reconnect
      socket.on('references:keys:updated', (keys, allDocs) => {
-        setReferenceKeys(oldKeys =>
-          allDocs ? new Set(keys) : new Set([...oldKeys, ...keys])
-        )
+        if (clientSideReferences) {
+          indexAllReferences(false)
+        } else {
+          setReferenceKeys(oldDocs =>
+            allDocs ? new Set(keys) : new Set([...oldDocs, ...keys])
+          )
+        }
      })
      indexAllReferences(false)
    }
@@ -127,14 +193,26 @@ export const ReferencesProvider: FC<React.PropsWithChildren> = ({
    return () => {
      eventEmitter.off('project:joined', handleProjectJoined)
    }
-  }, [eventEmitter, indexAllReferences, socket])
+  }, [eventEmitter, indexAllReferences, socket, clientSideReferences])
+
+  const searchLocalReferences = useCallback(
+    async (query: string): Promise<AdvancedReferenceSearchResult> => {
+      if (!indexerRef.current) {
+        return { hits: [] }
+      }
+      const indexer = await indexerRef.current
+      return await indexer.search(query)
+    },
+    []
+  )

  const value = useMemo(
    () => ({
      referenceKeys,
      indexAllReferences,
+      searchLocalReferences,
    }),
-    [indexAllReferences, referenceKeys]
+    [indexAllReferences, referenceKeys, searchLocalReferences]
  )

  return (
@@ -0,0 +1,24 @@
+import { ReferenceIndex } from './reference-index'
+import { Changes } from './types'
+
+export default class BasicReferenceIndex extends ReferenceIndex {
+  fileIndex: Map<string, Set<string>> = new Map()
+
+  updateIndex({ updates, deletes }: Changes): Set<string> {
+    for (const path of deletes) {
+      this.fileIndex.delete(path)
+    }
+    for (const { path, content } of updates) {
+      const fileReferences: Set<string> = new Set()
+      const entries = this.parseEntries(content)
+      for (const entry of entries) {
+        fileReferences.add(entry.EntryKey)
+      }
+      this.fileIndex.set(path, fileReferences)
+    }
+    this.keys = new Set(
+      this.fileIndex.values().flatMap(entry => Array.from(entry))
+    )
+    return this.keys
+  }
+}
@@ -0,0 +1,43 @@
+import Bib2Json from './bib2json'
+import { AdvancedReferenceSearchResult, Bib2JsonEntry, Changes } from './types'
+
+export abstract class ReferenceIndex {
+  keys: Set<string> = new Set()
+
+  abstract updateIndex({ updates, deletes }: Changes): void
+  async search(_query: string): Promise<AdvancedReferenceSearchResult> {
+    return { hits: [] }
+  }
+
+  getKeys(): Set<string> {
+    return this.keys
+  }
+
+  parseEntries(content: string): Bib2JsonEntry[] {
+    const allowedFields = ['author', 'journal', 'title', 'year', 'date']
+    // @ts-expect-error Bib2Json works as both a constructor and a function
+    const { entries } = Bib2Json(content, allowedFields)
+    for (const entry of entries) {
+      if (entry.Fields?.year) {
+        entry.Fields.year = parseInt(entry.Fields.year).toString()
+        if (entry.Fields.year === 'NaN') {
+          delete entry.Fields.year
+        }
+      }
+      setDefaultFields(entry.Fields)
+    }
+    return entries
+  }
+}
+
+function setDefaultFields(
+  fields: Partial<Bib2JsonEntry['Fields']>
+): Bib2JsonEntry['Fields'] {
+  const requiredFields = ['author', 'journal', 'title', 'date', 'year'] as const
+  for (const field of requiredFields) {
+    if (!fields[field]) {
+      fields[field] = ''
+    }
+  }
+  return fields as Bib2JsonEntry['Fields']
+}
@@ -0,0 +1,137 @@
+import { ProjectSnapshot } from '@/infrastructure/project-snapshot'
+import { generateSHA1Hash } from '@/shared/utils/sha1'
+import { AdvancedReferenceSearchResult, Changes } from './types'
+import { debugConsole } from '@/utils/debugging'
+import type { ReferenceWorkerResponse } from './references.worker'
+
+const ONE_MB = 1024 * 1024
+const MAX_BIB_DATA_SIZE = 6 * ONE_MB
+
+export class ReferenceIndexer {
+  private fileIndexHash: Map<string, string> = new Map()
+  private worker: Worker
+  private updateResolve: ((result: Set<string>) => void) | null = null
+  private searchResolve:
+    | ((result: AdvancedReferenceSearchResult) => void)
+    | null = null
+
+  constructor() {
+    this.worker = new Worker(
+      /* webpackChunkName: "references-worker" */
+      new URL('./references.worker.ts', import.meta.url),
+      { type: 'module' }
+    )
+    this.worker.addEventListener('message', evt => this.handleMessage(evt))
+  }
+
+  private handleMessage(event: MessageEvent) {
+    const data = event.data as ReferenceWorkerResponse
+    if (data.type === 'searchResult' && this.searchResolve) {
+      this.searchResolve(data.result)
+      this.searchResolve = null
+    } else if (data.type === 'updateKeys' && this.updateResolve) {
+      this.updateResolve(data.keys)
+      this.updateResolve = null
+    } else {
+      debugConsole.warn('Received unknown message from worker:', data.type)
+    }
+  }
+
+  async updateFromSnapshot(
+    snapshot: Pick<
+      ProjectSnapshot,
+      | 'getDocPaths'
+      | 'getDocContents'
+      | 'getBinaryFilePathsWithHash'
+      | 'getBinaryFileContents'
+    >,
+    {
+      dataLimit = MAX_BIB_DATA_SIZE,
+      signal,
+    }: { dataLimit?: number; signal: AbortSignal }
+  ): Promise<Set<string>> {
+    const nextFileHashIndex = new Map(this.fileIndexHash)
+    const previousPaths = new Set(this.fileIndexHash.keys())
+    let dataBudget = dataLimit
+    const docs = snapshot
+      .getDocPaths()
+      .filter(path => path.toLowerCase().endsWith('.bib'))
+
+    const changes: Changes = { updates: [], deletes: [] }
+    for (const path of docs) {
+      previousPaths.delete(path)
+      if (dataBudget <= 0) {
+        continue
+      }
+      const content = snapshot.getDocContents(path)?.slice(0, dataBudget)
+      if (content == null) {
+        continue
+      }
+      dataBudget -= content.length
+      const hash = generateSHA1Hash(content)
+      const possibleMatch = nextFileHashIndex.get(path)
+      if (possibleMatch === undefined || possibleMatch !== hash) {
+        // New or changed file
+        nextFileHashIndex.set(path, hash)
+        changes.updates.push({ path, content })
+      }
+    }
+
+    const files = snapshot
+      .getBinaryFilePathsWithHash()
+      .filter(({ path }) => path.toLowerCase().endsWith('.bib'))
+      .sort((a, b) => a.size - b.size)
+
+    for (const { path, hash, size } of files) {
+      if (signal.aborted) {
+        debugConsole.warn('Aborted indexing references due to signal')
+        return new Set()
+      }
+
+      previousPaths.delete(path)
+      if (nextFileHashIndex.get(path) === hash) {
+        dataBudget -= size
+        // Already indexed
+        continue
+      }
+      if (dataBudget <= 0) {
+        continue
+      }
+      const content = await snapshot.getBinaryFileContents(path, {
+        maxSize: dataBudget,
+      })
+      dataBudget -= content.length
+      nextFileHashIndex.set(path, hash)
+      changes.updates.push({ path, content })
+    }
+
+    previousPaths.forEach(path => {
+      // Deleted file
+      changes.deletes.push(path)
+      nextFileHashIndex.delete(path)
+    })
+
+    if (dataBudget <= 0) {
+      debugConsole.warn('Data budget exceeded while updating references index')
+    }
+
+    this.fileIndexHash = nextFileHashIndex
+
+    this.worker.postMessage({
+      type: 'update',
+      changes,
+    })
+
+    return new Promise(resolve => {
+      this.updateResolve = resolve
+    })
+  }
+
+  async search(query: string): Promise<AdvancedReferenceSearchResult> {
+    this.worker.postMessage({ type: 'search', query })
+    const { promise, resolve } =
+      Promise.withResolvers<AdvancedReferenceSearchResult>()
+    this.searchResolve = resolve
+    return promise
+  }
+}
@@ -0,0 +1,47 @@
+import BasicReferenceIndex from './basic-reference-index'
+import { ReferenceIndex } from './reference-index'
+import { AdvancedReferenceSearchResult, Changes } from './types'
+import importOverleafModules from '../../../../macros/import-overleaf-module.macro'
+
+interface IndexConstructor {
+  new (): ReferenceIndex
+}
+
+const indices = importOverleafModules('referenceIndices') as {
+  import: { default: IndexConstructor }
+  path: string
+}[]
+
+export type ReferenceWorkerRequest =
+  | { type: 'update'; changes: Changes }
+  | { type: 'search'; query: string }
+
+export type ReferenceWorkerResponse =
+  | { type: 'updateKeys'; keys: Set<string> }
+  | { type: 'searchResult'; result: AdvancedReferenceSearchResult }
+
+function createIndex(): ReferenceIndex {
+  const Klass = indices[0]?.import.default ?? BasicReferenceIndex
+  return new Klass()
+}
+
+const indexer: ReferenceIndex = createIndex()
+
+self.addEventListener('message', async (event: MessageEvent) => {
+  const message = event.data as ReferenceWorkerRequest
+  switch (message.type) {
+    case 'update':
+      indexer.updateIndex(message.changes)
+      self.postMessage({ type: 'updateKeys', keys: indexer.getKeys() })
+      break
+
+    case 'search': {
+      const result = await indexer.search(message.query)
+      self.postMessage({ type: 'searchResult', result })
+      break
+    }
+
+    default:
+      console.error('Unknown message type:', message)
+  }
+})
@@ -0,0 +1,23 @@
+export type Bib2JsonEntry = {
+  EntryKey: string
+  Fields: {
+    author: string
+    date: string
+    journal: string
+    title: string
+    year: string
+  }
+}
+
+export type AdvancedReferenceSearchResult = {
+  hits: {
+    _source: Bib2JsonEntry
+  }[]
+}
+
+export type ReferenceEntry = Map<string, Bib2JsonEntry>
+
+export type Changes = {
+  updates: { path: string; content: string }[]
+  deletes: string[]
+}