fix(attachments): abort embedding on cancel and cap oversized document

Author: kfaracik

__tests__/ChatBar.test.tsx

@@ -1,6 +1,8 @@
 import React from 'react';
 import { render, screen, fireEvent, act } from '@testing-library/react-native';
 import type { LLMStore } from '../store/llmStore';
+import type { Attachment } from '../hooks/useAttachment';
+import type { PermissionStatus } from 'react-native-audio-api';
 
 // ── mocks ─────────────────────────────────────────────────────────────────────
 
@@ -27,7 +29,7 @@ jest.mock('../store/llmStore', () => ({
 }));
 
 const mockUseAttachment = {
-  attachments: [] as any[],
+  attachments: [] as Attachment[],
   sheetRef: { current: null },
   pickFromLibrary: jest.fn(),
   pickFromCamera: jest.fn(),
@@ -49,7 +51,12 @@ jest.mock('../components/bottomSheets/AttachmentSheet', () => {
     onPickFromCamera,
     onPickDocument,
     isVisionModel,
-  }: any) => (
+  }: {
+    onPickFromLibrary: () => void;
+    onPickFromCamera: () => void;
+    onPickDocument: () => void;
+    isVisionModel: boolean;
+  }) => (
     <View testID="attachment-sheet">
       {isVisionModel && (
         <>
@@ -73,7 +80,13 @@ jest.mock('../components/bottomSheets/AttachmentSheet', () => {
 
 jest.mock('../components/chat-screen/AttachmentThumbnail', () => {
   const { View, TouchableOpacity, Text } = require('react-native');
-  return ({ attachment, onRemove }: any) => (
+  return ({
+    attachment,
+    onRemove,
+  }: {
+    attachment: Attachment;
+    onRemove: () => void;
+  }) => (
     <View testID={`attachment-thumb-${attachment.id}`}>
       <Text>{attachment.name || attachment.uri}</Text>
       <TouchableOpacity
@@ -87,8 +100,14 @@ jest.mock('../components/chat-screen/AttachmentThumbnail', () => {
 });
 
 jest.mock('../components/chat-screen/ChatSpeechInput', () => {
-  const { View, TouchableOpacity, Text } = require('react-native');
-  return ({ onSubmit, onCancel }: any) => (
+  const { View, TouchableOpacity } = require('react-native');
+  return ({
+    onSubmit,
+    onCancel,
+  }: {
+    onSubmit: (transcript: string) => void;
+    onCancel: () => void;
+  }) => (
     <View testID="speech-input">
       <TouchableOpacity
         testID="speech-submit"
@@ -101,7 +120,7 @@ jest.mock('../components/chat-screen/ChatSpeechInput', () => {
 
 jest.mock('../components/chat-screen/PromptSuggestions', () => {
   const { TouchableOpacity, Text } = require('react-native');
-  return ({ onSelectPrompt }: any) => (
+  return ({ onSelectPrompt }: { onSelectPrompt: (prompt: string) => void }) => (
     <TouchableOpacity
       testID="prompt-suggestion"
       onPress={() => onSelectPrompt('Suggested prompt')}
@@ -124,7 +143,18 @@ jest.mock('../components/chat-screen/ChatBarActions', () => {
     onThinkingToggle,
     thinkingEnabled,
     onAttach,
-  }: any) => (
+  }: {
+    userInput: string;
+    hasAttachments: boolean;
+    onSend: () => void;
+    isGenerating: boolean;
+    isProcessingPrompt: boolean;
+    onInterrupt: () => void;
+    onSpeechInput: () => void;
+    onThinkingToggle: () => void;
+    thinkingEnabled: boolean;
+    onAttach: () => void;
+  }) => (
     <View testID="chat-bar-actions">
       <TouchableOpacity testID="attach-btn" onPress={onAttach}>
         <Text>+</Text>
@@ -185,7 +215,7 @@ const defaultProps = {
   onSelectModel: jest.fn(),
   onSelectPrompt: jest.fn(),
   model: downloadedModel,
-  scrollRef: { current: null } as any,
+  scrollRef: { current: null },
   isAtBottom: true,
   isVisionModel: false,
   thinkingEnabled: false,
@@ -218,7 +248,7 @@ beforeEach(() => {
   jest.spyOn(console, 'warn').mockImplementation(() => {});
   // Default: permission granted
   mockAudioManager.requestRecordingPermissions.mockResolvedValue(
-    'Granted' as any
+    'Granted' as PermissionStatus
   );
 });
 
@@ -357,7 +387,7 @@ describe('speech input', () => {
 
   it('shows toast and stays in text mode when microphone permission is denied', async () => {
     mockAudioManager.requestRecordingPermissions.mockResolvedValue(
-      'Denied' as any
+      'Denied' as PermissionStatus
     );
     renderBar();
     await act(async () => {
@@ -430,7 +460,7 @@ describe('speech input', () => {
     // Override speech mock to submit empty string
     jest.mock('../components/chat-screen/ChatSpeechInput', () => {
       const { View, TouchableOpacity } = require('react-native');
-      return ({ onSubmit }: any) => (
+      return ({ onSubmit }: { onSubmit: (transcript: string) => void }) => (
         <View testID="speech-input">
           <TouchableOpacity
             testID="speech-submit"

__tests__/hybridRetrieval.test.ts

@@ -1,5 +1,6 @@
 import { hybridRetrieve } from '../utils/hybridRetrieval';
 import * as keywordIndex from '../database/keywordIndex';
+import type { OPSQLiteVectorStore } from '@react-native-rag/op-sqlite';
 
 jest.mock('../database/keywordIndex', () => ({
   keywordSearch: jest.fn(),
@@ -8,8 +9,8 @@ jest.mock('../database/keywordIndex', () => ({
 const mockKeywordSearch = keywordIndex.keywordSearch as jest.Mock;
 
 const makeVectorStore = (
-  queryResults: any[],
-  vectorsById: Record<string, any>
+  queryResults: unknown[],
+  vectorsById: Record<string, unknown>
 ) =>
   ({
     query: jest.fn().mockResolvedValue(queryResults),
@@ -20,7 +21,7 @@ const makeVectorStore = (
           rows: ids.map((id) => vectorsById[id]).filter(Boolean),
         })),
     },
-  }) as any;
+  }) as unknown as OPSQLiteVectorStore;
 
 describe('hybridRetrieve', () => {
   beforeEach(() => {

__tests__/llmStore.test.ts

@@ -1,6 +1,9 @@
 import { useLLMStore } from '../store/llmStore';
 import { LLMModule } from 'react-native-executorch';
 import * as chatRepository from '../database/chatRepository';
+import type { Message } from '../database/chatRepository';
+import type { Model } from '../database/modelRepository';
+import type { SQLiteDatabase } from 'expo-sqlite';
 import * as Feedback from '../utils/Feedback';
 import { prepareMessagesForLLM } from '../utils/promptUtils';
 
@@ -29,7 +32,7 @@ const noSources = async () => ({
   preferredSourceDocuments: [],
 });
 
-const mockDb = {} as any;
+const mockDb = {} as unknown as SQLiteDatabase;
 
 const baseModel = {
   id: 1,
@@ -64,7 +67,7 @@ beforeEach(() => {
   mockLLMModule.fromModelName.mockImplementation(
     async (_namedSources, _onProgress, onToken) => {
       capturedTokenCallback = onToken;
-      return mockInstance as any;
+      return mockInstance as unknown as LLMModule;
     }
   );
 
@@ -89,7 +92,7 @@ beforeEach(() => {
   mockLLMModule.fromModelName.mockImplementation(
     async (_namedSources, _onProgress, onToken) => {
       capturedTokenCallback = onToken;
-      return mockInstance as any;
+      return mockInstance as unknown as LLMModule;
     }
   );
 });
@@ -112,7 +115,7 @@ describe('loadModel', () => {
     mockLLMModule.fromModelName.mockImplementation(async (...args) => {
       wasLoading = useLLMStore.getState().isLoading;
       capturedTokenCallback = args[4];
-      return mockInstance as any;
+      return mockInstance as unknown as LLMModule;
     });
 
     await useLLMStore.getState().loadModel(baseModel);
@@ -142,7 +145,7 @@ describe('loadModel', () => {
     mockInstance = makeMockInstance();
     mockLLMModule.fromModelName.mockImplementation(async (...args) => {
       capturedTokenCallback = args[4];
-      return mockInstance as any;
+      return mockInstance as unknown as LLMModule;
     });
     await useLLMStore.getState().loadModel({ ...baseModel, id: 2 });
 
@@ -367,7 +370,7 @@ describe('sendChatMessage', () => {
   });
 
   it('adds user message and assistant placeholder to activeChatMessages before generating', async () => {
-    let messagesBeforeGenerate: any[] = [];
+    let messagesBeforeGenerate: Message[] = [];
     mockInstance.generate.mockImplementation(async () => {
       messagesBeforeGenerate = useLLMStore.getState().activeChatMessages;
       return 'response';
@@ -536,7 +539,7 @@ describe('sendChatMessage imagePath', () => {
         modelName: 'LFM VL',
         vision: true,
         featured: true,
-      } as any,
+      } as Model,
       activeChatId: 1,
       activeChatMessages: [],
     });

__tests__/sourceStore.test.ts

@@ -6,6 +6,7 @@ import { useLLMStore } from '../store/llmStore';
 import type { SQLiteDatabase } from 'expo-sqlite';
 import type { OPSQLiteVectorStore } from '@react-native-rag/op-sqlite';
 import type { LFMEmbeddings } from '../utils/lfmEmbeddings';
+import { MAX_SOURCE_CHUNKS } from '../constants/retrieval';
 
 jest.mock('../database/sourcesRepository');
 jest.mock('../utils/fileReaders');
@@ -25,8 +26,10 @@ jest.mock('@react-native-rag/op-sqlite', () => ({}));
 
 const mockDb = {} as Partial<SQLiteDatabase> as SQLiteDatabase;
 const vectorStoreAdd = jest.fn();
+const vectorStoreDelete = jest.fn();
 const mockVectorStore = {
   add: vectorStoreAdd,
+  delete: vectorStoreDelete,
 } as Partial<OPSQLiteVectorStore> as OPSQLiteVectorStore;
 
 const mockReadDocumentText = fileReaders.readDocumentText as jest.Mock;
@@ -141,12 +144,58 @@ describe('addSource', () => {
       .addSource(baseSource, '/path/doc.txt', mockVectorStore);
 
     const sources = useSourceStore.getState().sources;
-    expect(result).toEqual({ success: true, sourceId: 99 });
+    expect(result).toEqual({ success: true, sourceId: 99, truncated: false });
     expect(sources).toHaveLength(1);
     expect(sources[0].id).toBe(99);
     expect(sources[0].isProcessing).toBe(false);
   });
 
+  it('caps embedded chunks at MAX_SOURCE_CHUNKS and flags the result truncated', async () => {
+    mockReadDocumentText.mockResolvedValue('content');
+    mockInsertSource.mockResolvedValue(99);
+    const manyChunks = Array.from(
+      { length: MAX_SOURCE_CHUNKS + 1 },
+      (_, i) => `chunk-${i}`
+    );
+    MockSplitter.mockImplementation(() => ({
+      splitText: jest.fn().mockResolvedValue(manyChunks),
+    }));
+
+    const result = await useSourceStore
+      .getState()
+      .addSource(baseSource, '/path/doc.txt', mockVectorStore);
+
+    expect(result).toEqual({ success: true, sourceId: 99, truncated: true });
+    expect(vectorStoreAdd).toHaveBeenCalledTimes(MAX_SOURCE_CHUNKS);
+  });
+
+  it('aborts embedding and rolls back the partial source when the signal is aborted', async () => {
+    mockReadDocumentText.mockResolvedValue('content');
+    mockInsertSource.mockResolvedValue(99);
+    MockSplitter.mockImplementation(() => ({
+      splitText: jest.fn().mockResolvedValue(['chunk-a', 'chunk-b']),
+    }));
+    const controller = new AbortController();
+    controller.abort();
+
+    const result = await useSourceStore
+      .getState()
+      .addSource(
+        baseSource,
+        '/path/doc.txt',
+        mockVectorStore,
+        undefined,
+        undefined,
+        controller.signal
+      );
+
+    expect(result).toEqual({ success: false, cancelled: true });
+    expect(vectorStoreAdd).not.toHaveBeenCalled();
+    expect(vectorStoreDelete).toHaveBeenCalledTimes(1);
+    expect(mockDeleteSource).toHaveBeenCalledWith(mockDb, 99);
+    expect(useSourceStore.getState().sources).toHaveLength(0);
+  });
+
   it('passes firstChunk to insertSource', async () => {
     mockReadDocumentText.mockResolvedValue('content');
     mockInsertSource.mockResolvedValue(1);
@@ -316,10 +365,10 @@ describe('cleanupOrphanedSources', () => {
     mockGetOrphanedSources.mockResolvedValue(orphaned);
     mockDeleteSource.mockResolvedValue(undefined);
 
-    const vectorStoreDelete = jest.fn();
+    const orphanVectorStoreDelete = jest.fn();
     const mockVectorStoreWithDelete = {
       add: vectorStoreAdd,
-      delete: vectorStoreDelete,
+      delete: orphanVectorStoreDelete,
     } as Partial<OPSQLiteVectorStore> as OPSQLiteVectorStore;
 
     await useSourceStore
@@ -328,7 +377,7 @@ describe('cleanupOrphanedSources', () => {
 
     expect(mockGetOrphanedSources).toHaveBeenCalledWith(mockDb);
     expect(mockDeleteSource).toHaveBeenCalledWith(mockDb, 5);
-    expect(vectorStoreDelete).toHaveBeenCalledWith({
+    expect(orphanVectorStoreDelete).toHaveBeenCalledWith({
       predicate: expect.any(Function),
     });
   });

constants/retrieval.ts

@@ -41,6 +41,9 @@ export const ANSWER_CITATION_OVERLAP_RATIO = 0.5;
 export const TEXT_SPLITTER_CHUNK_SIZE = 1000;
 export const TEXT_SPLITTER_CHUNK_OVERLAP = 200;
 
+/** Safety backstop on chunks embedded per source; high enough that large real documents (~1.6 MB of text) index in full, low enough to stop a pathological multi-MB file from embedding for tens of minutes. */
+export const MAX_SOURCE_CHUNKS = 2000;
+
 /** Min matched run to treat as overlap when stitching passages — below the splitter overlap, above coincidental repetition. */
 export const MIN_STITCH_OVERLAP = 24;