fix(i18n): retry failed locale loads (#1206)

* fix(i18n): retry failed locale loads

* test(i18n): reset locale cache in retry test

---------

Co-authored-by: Isaac Sanchez <isanchez-hawkins@arize.com>

packages/ui/src/lib/i18n/store.test.ts

@@ -0,0 +1,45 @@
+import { beforeEach, describe, expect, test } from 'bun:test';
+import { DEFAULT_LOCALE, type Locale } from './runtime';
+import { resetI18nDictionaryCacheForTests, useI18nStore } from './store';
+
+const defaultDictionary = useI18nStore.getState().dictionary;
+
+const resetStore = () => {
+  resetI18nDictionaryCacheForTests();
+  useI18nStore.setState({
+    locale: DEFAULT_LOCALE,
+    dictionary: defaultDictionary,
+    loadingLocale: null,
+  });
+};
+
+const waitForLocaleLoadToSettle = async (locale: Locale) => {
+  for (let attempt = 0; attempt < 20; attempt += 1) {
+    if (useI18nStore.getState().loadingLocale !== locale) {
+      return;
+    }
+    await new Promise((resolve) => setTimeout(resolve, 0));
+  }
+  throw new Error(`Timed out waiting for ${locale} dictionary load`);
+};
+
+describe('i18n store', () => {
+  beforeEach(resetStore);
+
+  test('retries loading the active locale when it is not cached', async () => {
+    useI18nStore.setState({
+      locale: 'es',
+      dictionary: defaultDictionary,
+      loadingLocale: null,
+    });
+
+    try {
+      useI18nStore.getState().setLocale('es');
+
+      expect(useI18nStore.getState().loadingLocale).toBe('es');
+      await waitForLocaleLoadToSettle('es');
+    } finally {
+      resetStore();
+    }
+  });
+});

packages/ui/src/lib/i18n/store.ts

@@ -15,6 +15,11 @@ type I18nState = {
 
 const dictionaries = new Map<Locale, I18nDictionary>([[DEFAULT_LOCALE, enDict]]);
 
+export function resetI18nDictionaryCacheForTests(): void {
+  dictionaries.clear();
+  dictionaries.set(DEFAULT_LOCALE, enDict);
+}
+
 async function loadDictionary(locale: Locale): Promise<I18nDictionary> {
   const cached = dictionaries.get(locale);
   if (cached) {
@@ -44,13 +49,13 @@ export const useI18nStore = create<I18nState>()((set, get) => ({
   loadingLocale: null,
   setLocale: (locale) => {
     const current = get();
-    if (current.locale === locale && current.loadingLocale !== locale) {
+    const cached = dictionaries.get(locale);
+    if (current.locale === locale && current.loadingLocale !== locale && cached) {
       return;
     }
 
     writeStoredLocale(locale);
 
-    const cached = dictionaries.get(locale);
     set({
       locale,
       dictionary: cached ?? current.dictionary,