fix(chat): don't flash the legacy notice for in-flight messages

kfaracik · kfaracik · commit a0d3cc2ec0c7 · 2026-07-31T13:06:11.000+02:00
diff --git a/__tests__/ChatBar.test.tsx b/__tests__/ChatBar.test.tsx
@@ -6,6 +6,10 @@ import type { PermissionStatus } from 'react-native-audio-api';
 
 // ── mocks ─────────────────────────────────────────────────────────────────────
 
+const mockRunWithModelOffloaded = jest.fn(
+  async (operation: () => Promise<unknown>) => operation()
+);
+
 jest.mock('../context/ThemeContext', () => ({
   useTheme: () => ({
     theme: {
@@ -23,6 +27,8 @@ jest.mock('../store/llmStore', () => ({
       interrupt: jest.fn(),
       loadModel: jest.fn(),
       model: null,
+      runWithModelOffloaded:
+        mockRunWithModelOffloaded as unknown as LLMStore['runWithModelOffloaded'],
     };
     return selector ? selector(state) : state;
   }),
@@ -229,6 +235,8 @@ beforeEach(() => {
         interrupt: jest.fn(),
         loadModel: jest.fn(),
         model: null,
+        runWithModelOffloaded:
+          mockRunWithModelOffloaded as unknown as LLMStore['runWithModelOffloaded'],
       };
       return selector ? selector(state) : state;
     }
@@ -237,6 +245,7 @@ beforeEach(() => {
   mockUseAttachment.openSheet.mockClear();
   mockUseAttachment.clearAll.mockClear();
   mockUseAttachment.removeAttachment.mockClear();
+  mockRunWithModelOffloaded.mockClear();
   jest.clearAllMocks();
   jest.spyOn(console, 'error').mockImplementation(() => {});
   jest.spyOn(console, 'warn').mockImplementation(() => {});
@@ -483,9 +492,15 @@ describe('attachment', () => {
     expect(screen.getByTestId('attach-btn')).toBeTruthy();
   });
 
-  it('opens attachment sheet when + button is pressed', () => {
+  it('offloads the LLM before opening the attachment sheet', async () => {
     renderBar();
-    fireEvent.press(screen.getByTestId('attach-btn'));
+    await act(async () => {
+      fireEvent.press(screen.getByTestId('attach-btn'));
+    });
+    expect(mockRunWithModelOffloaded).toHaveBeenCalledWith(
+      expect.any(Function),
+      { restore: false }
+    );
     expect(mockUseAttachment.openSheet).toHaveBeenCalled();
   });
 
diff --git a/__tests__/VectorStoreContext.test.tsx b/__tests__/VectorStoreContext.test.tsx
@@ -7,6 +7,9 @@ const mockStoreInstances: Array<{
   load: jest.Mock;
   unload: jest.Mock;
 }> = [];
+const mockEmbeddingInstances: Array<{
+  unload: jest.Mock;
+}> = [];
 
 jest.mock('@react-native-rag/op-sqlite', () => ({
   OPSQLiteVectorStore: jest
@@ -41,7 +44,13 @@ jest.mock('../utils/embeddingModel', () => ({
 }));
 
 jest.mock('../utils/lfmEmbeddings', () => ({
-  LFMEmbeddings: jest.fn().mockImplementation(() => ({})),
+  LFMEmbeddings: jest.fn().mockImplementation(() => {
+    const instance = {
+      unload: jest.fn().mockResolvedValue(undefined),
+    };
+    mockEmbeddingInstances.push(instance);
+    return instance;
+  }),
 }));
 
 jest.mock('../store/embeddingModelStore', () => ({
@@ -80,6 +89,7 @@ const renderProvider = () =>
 beforeEach(() => {
   jest.clearAllMocks();
   mockStoreInstances.length = 0;
+  mockEmbeddingInstances.length = 0;
   mockMigrate.mockResolvedValue(undefined);
 });
 
@@ -105,6 +115,14 @@ describe('VectorStoreProvider init/teardown chain', () => {
     expect(mockStoreInstances[1].unload).not.toHaveBeenCalled();
   });
 
+  it('releases the downloaded embedding model after initialization', async () => {
+    renderProvider();
+    await flush();
+
+    expect(mockEmbeddingInstances).toHaveLength(1);
+    expect(mockEmbeddingInstances[0].unload).toHaveBeenCalledTimes(1);
+  });
+
   it('unmount during init unloads the store exactly once', async () => {
     const deferred = createDeferred<void>();
     mockMigrate.mockReturnValueOnce(deferred.promise);
diff --git a/__tests__/lfmEmbeddings.test.ts b/__tests__/lfmEmbeddings.test.ts
@@ -8,6 +8,76 @@ const makeEmbeddings = () =>
     tokenizerSource: 'file://tokenizer.json',
   });
 
+describe('LFMEmbeddings.runWithLoadedModel', () => {
+  it('loads for an operation and always unloads afterwards', async () => {
+    const embeddings = makeEmbeddings();
+    const load = jest.spyOn(embeddings, 'load').mockResolvedValue(embeddings);
+    const unload = jest.spyOn(embeddings, 'unload').mockResolvedValue();
+    const operation = jest.fn().mockResolvedValue('result');
+
+    await expect(embeddings.runWithLoadedModel(operation)).resolves.toBe(
+      'result'
+    );
+    expect(load).toHaveBeenCalledTimes(1);
+    expect(operation).toHaveBeenCalledTimes(1);
+    expect(unload).toHaveBeenCalledTimes(1);
+    expect(load.mock.invocationCallOrder[0]).toBeLessThan(
+      operation.mock.invocationCallOrder[0]
+    );
+    expect(operation.mock.invocationCallOrder[0]).toBeLessThan(
+      unload.mock.invocationCallOrder[0]
+    );
+  });
+
+  it('unloads when the operation fails', async () => {
+    const embeddings = makeEmbeddings();
+    jest.spyOn(embeddings, 'load').mockResolvedValue(embeddings);
+    const unload = jest.spyOn(embeddings, 'unload').mockResolvedValue();
+
+    await expect(
+      embeddings.runWithLoadedModel(async () => {
+        throw new Error('embedding failed');
+      })
+    ).rejects.toThrow('embedding failed');
+    expect(unload).toHaveBeenCalledTimes(1);
+  });
+
+  it('serializes operations so one cannot unload another model session', async () => {
+    const embeddings = makeEmbeddings();
+    const load = jest.spyOn(embeddings, 'load').mockResolvedValue(embeddings);
+    const unload = jest.spyOn(embeddings, 'unload').mockResolvedValue();
+    let finishFirst!: () => void;
+    let markFirstStarted!: () => void;
+    const firstGate = new Promise<void>((resolve) => {
+      finishFirst = resolve;
+    });
+    const firstStarted = new Promise<void>((resolve) => {
+      markFirstStarted = resolve;
+    });
+    const order: string[] = [];
+
+    const first = embeddings.runWithLoadedModel(async () => {
+      order.push('first-start');
+      markFirstStarted();
+      await firstGate;
+      order.push('first-end');
+    });
+    const second = embeddings.runWithLoadedModel(async () => {
+      order.push('second');
+    });
+
+    await firstStarted;
+    expect(order).toEqual(['first-start']);
+
+    finishFirst();
+    await Promise.all([first, second]);
+
+    expect(order).toEqual(['first-start', 'first-end', 'second']);
+    expect(load).toHaveBeenCalledTimes(2);
+    expect(unload).toHaveBeenCalledTimes(2);
+  });
+});
+
 describe('embedding input limits', () => {
   it('sends a short query through with its prefix intact', async () => {
     const embeddings = makeEmbeddings();
diff --git a/__tests__/useAttachment.test.ts b/__tests__/useAttachment.test.ts
@@ -30,13 +30,16 @@ import { launchImageLibrary, launchCamera } from 'react-native-image-picker';
 import * as DocumentPicker from 'expo-document-picker';
 import { useAttachment } from '../hooks/useAttachment';
 import { useEmbeddingModelStore } from '../store/embeddingModelStore';
+import { useVectorStore } from '../context/VectorStoreContext';
 
 const mockLaunchImageLibrary = launchImageLibrary as jest.Mock;
 const mockLaunchCamera = launchCamera as jest.Mock;
 const mockGetDocumentAsync = DocumentPicker.getDocumentAsync as jest.Mock;
+const mockUseVectorStore = useVectorStore as jest.Mock;
 
 beforeEach(() => {
   jest.clearAllMocks();
+  mockUseVectorStore.mockReturnValue({ vectorStore: {}, embeddings: null });
   useEmbeddingModelStore.setState({ status: 'ready', progress: 1 });
 });
 
@@ -99,6 +102,13 @@ describe('useAttachment', () => {
     const mockAddSource = jest
       .fn()
       .mockResolvedValue({ success: true, sourceId: 42 });
+    const runWithLoadedModel = jest.fn(
+      async (operation: () => Promise<unknown>) => operation()
+    );
+    mockUseVectorStore.mockReturnValue({
+      vectorStore: {},
+      embeddings: { runWithLoadedModel },
+    });
     const { useSourceStore } = require('../store/sourceStore');
     useSourceStore.getState.mockReturnValue({
       addSource: mockAddSource,
@@ -115,6 +125,7 @@ describe('useAttachment', () => {
     expect(att.type).toBe('document');
     expect(att.sourceId).toBe(42);
     expect(att.status).toBe('ready');
+    expect(runWithLoadedModel).toHaveBeenCalledTimes(1);
   });
 
   describe('abandoned source cleanup', () => {
diff --git a/components/chat-screen/ChatBar.tsx b/components/chat-screen/ChatBar.tsx
@@ -183,18 +183,24 @@ const ChatBar = ({
     interrupt,
     loadModel,
     model: loadedModel,
+    runWithModelOffloaded,
   } = useLLMStore();
   const loadSelectedModel = useCallback(async () => {
     if (model?.isDownloaded && loadedModel?.id !== model.id) {
       return loadModel(model);
     }
   }, [model, loadedModel, loadModel]);
 
-  const handleAttach = useCallback(() => {
+  const handleAttach = useCallback(async () => {
     Keyboard.dismiss();
-    loadSelectedModel();
-    openSheet();
-  }, [loadSelectedModel, openSheet]);
+    try {
+      await runWithModelOffloaded(async () => {}, { restore: false });
+    } catch (error) {
+      console.error('Failed to offload model before attachment picker:', error);
+    } finally {
+      openSheet();
+    }
+  }, [openSheet, runWithModelOffloaded]);
 
   const imageAttachment = attachments.find((a) => a.type === 'image');
   const hasLoadingAttachment = attachments.some((a) => a.status === 'loading');
diff --git a/components/chat-screen/ChatScreen.tsx b/components/chat-screen/ChatScreen.tsx
@@ -89,6 +89,7 @@ export default function ChatScreen({
     sendChatMessage,
     loadModel,
     model: loadedModel,
+    runWithModelOffloaded,
     generationError,
     retryLastGeneration,
   } = useLLMStore();
@@ -252,15 +253,21 @@ export default function ChatScreen({
       let sourceDocuments: SourceDocument[] = [];
       let preferredSourceDocuments: SourceDocument[] = [];
       if (vectorStore) {
-        ({ context, sourceDocuments, preferredSourceDocuments } =
-          await buildMessageSources({
+        const prepareSources = () =>
+          buildMessageSources({
             userInput,
             attachmentSourceIds,
             enabledSources,
             sources: allSources,
             vectorStore,
             embeddings,
-          }));
+          });
+        ({ context, sourceDocuments, preferredSourceDocuments } = embeddings
+          ? await runWithModelOffloaded(
+              () => embeddings.runWithLoadedModel(prepareSources),
+              { restore: false }
+            )
+          : await prepareSources());
       }
 
       // Enable new sources for this chat (persists for future messages)
diff --git a/context/VectorStoreContext.tsx b/context/VectorStoreContext.tsx
@@ -11,6 +11,7 @@ import { LFMEmbeddings } from '../utils/lfmEmbeddings';
 import { ensureKeywordIndex } from '../database/keywordIndex';
 import { isEmbeddingModelDownloaded } from '../utils/embeddingModel';
 import { useEmbeddingModelStore } from '../store/embeddingModelStore';
+import { useLLMStore } from '../store/llmStore';
 
 const VectorStoreContext = createContext<{
   vectorStore: OPSQLiteVectorStore | null;
@@ -80,7 +81,13 @@ export const VectorStoreProvider = ({
 
         if (cancelled) return;
         if (downloaded) {
-          await store.load();
+          await useLLMStore.getState().runWithModelOffloaded(
+            async () => {
+              await store.load();
+              await lfmEmbeddings.unload();
+            },
+            { restore: false }
+          );
         }
 
         if (cancelled) return;
diff --git a/hooks/useAttachment.ts b/hooks/useAttachment.ts
@@ -7,6 +7,7 @@ import Toast from 'react-native-toast-message';
 import { useSourceStore } from '../store/sourceStore';
 import { useVectorStore } from '../context/VectorStoreContext';
 import { useEmbeddingModelStore } from '../store/embeddingModelStore';
+import { useLLMStore } from '../store/llmStore';
 import { documentErrorMessage } from '../utils/documentErrorMessage';
 
 export interface Attachment {
@@ -194,14 +195,25 @@ export const useAttachment = () => {
           prev.map((a) => (a.id === attachmentId ? { ...a, progress } : a))
         );
       };
-      const result = await addSource(
-        newSource,
-        asset.uri,
-        vectorStore!,
-        embeddings,
-        handleProgress,
-        abortController.signal
-      );
+      const addDocumentSource = () =>
+        addSource(
+          newSource,
+          asset.uri,
+          vectorStore!,
+          embeddings,
+          handleProgress,
+          abortController.signal
+        );
+      const indexDocumentSource = () =>
+        embeddings
+          ? useLLMStore
+              .getState()
+              .runWithModelOffloaded(
+                () => embeddings.runWithLoadedModel(addDocumentSource),
+                { restore: false }
+              )
+          : addDocumentSource();
+      const result = await indexDocumentSource();
       if (result.cancelled) return;
       const isCurrentDocumentRequest =
         attachmentRequestRef.current === requestId &&
@@ -271,9 +283,16 @@ export const useAttachment = () => {
 
   const downloadModelAndContinue = useCallback(async () => {
     if (!vectorStore) return;
-    const ready = await useEmbeddingModelStore
-      .getState()
-      .ensureReady(vectorStore);
+    const ready = await useLLMStore.getState().runWithModelOffloaded(
+      async () => {
+        const loaded = await useEmbeddingModelStore
+          .getState()
+          .ensureReady(vectorStore);
+        await embeddings?.unload();
+        return loaded;
+      },
+      { restore: false }
+    );
     if (!ready) {
       Toast.show({
         type: 'defaultToast',
@@ -285,7 +304,7 @@ export const useAttachment = () => {
       embeddingDownloadSheetRef.current?.dismiss();
       await runDocumentPicker();
     }
-  }, [vectorStore, runDocumentPicker]);
+  }, [vectorStore, embeddings, runDocumentPicker]);
 
   const removeAttachment = useCallback(
     (id: string) => {
diff --git a/utils/lfmEmbeddings.ts b/utils/lfmEmbeddings.ts
