Merge pull request #498 from off-grid-ai/fix/load-anyway-all-model-types

fix: offer "Load Anyway" memory override for all model types (image parity)

Author: alichherawalla

CLAUDE.md

@@ -134,6 +134,12 @@ Always write **both** unit tests and integration tests for new features and sign
 
 Do not consider a feature complete with only unit tests. Integration tests catch wiring bugs, incorrect data flow between layers, and lifecycle issues that unit tests miss.
 
+### Coverage (core AND pro, 100% on new code)
+
+**Coverage must include the `pro/` submodule, not just `src/`.** `jest.config.js` `collectCoverageFrom` collects from `pro/**` whenever the submodule is checked out (gated on `proExists`, so open-core CI without pro still runs). Pro features (TTS/audio, MCP, and other paid surfaces) are exercised by the pro-dependent suites in this repo's `__tests__/` — they must be *measured*, never invisible to coverage. Do not add code to `pro/` without its coverage counting.
+
+**All NEW code ships at 100% - statements, branches, functions, AND lines.** Every new branch/condition (each side of every `if`/ternary/`??`/`||`, each catch, each early return) needs a test that exercises it. Enforce it with a per-file `coverageThreshold` entry at 100 for each new standalone module (core or pro); for a *changed* legacy file, cover every new/changed branch even though the whole file isn't held to 100. New code with an uncovered branch is not done.
+
 **Use mocks very sparingly - a green suite must mean the real thing works, not that a mock returned what it was told.** Mock only what you genuinely cannot run in the test environment (native modules, the network, the device clock). Everything else - the service under test, the stores it writes, the logic across layers - runs for real. A test that mocks the very thing it is asserting (so it would pass even if the implementation were deleted) is worse than no test: it hides the broken behaviour behind a false green. Prefer driving the real class/store/reducer and asserting the observable outcome. When you must stub a boundary, keep the stub dumb (return plain data) and let the real logic on top of it do the work. If a behaviour can only be proven by mocking out the behaviour, that is the signal to test it at a higher layer (integration) or on-device (Provit) instead.
 
 **Design to SOLID with real abstraction layers (not incidental ones).** These are the same rules as the Architecture section above, restated as a standing expectation for every change: one responsibility per module (SRP); callers depend on an interface/service, never on a concrete implementation or a `kind===`/`instanceof`/`Platform.OS`-mechanism branch (DIP); a new implementation (engine, provider, backend) drops in behind the existing seam with zero caller changes (OCP); any implementation is substitutable through the interface (LSP); interfaces are segregated so an implementation never stubs methods it can't support. The abstraction layer must be a genuine owning seam - a service that owns the state machine, resources, and side-effects - not a thin pass-through that leaks the concretes upward. If a fix would add a second concrete branch in a caller, build/extend the seam instead.

__tests__/integration/audio/streamingStateMachine.test.ts

@@ -50,6 +50,7 @@ jest.mock('../../../pro/audio/ttsStore', () => ({
 import { useTTSStore } from '../../../pro/audio/ttsStore';
 import {
   feedStreamingText, finishStreamingText, resetStreamingSpeech, isStreamingSpeechActive, _setSpeakTimeoutForTest,
+  stopStreamingSpeechForTurn,
 } from '../../../pro/audio/streamingSpeech';
 import { _setSmSink, type SmEvent } from '../../../pro/audio/ttsLog';
 
@@ -218,3 +219,49 @@ describe('streaming state machine — reset always reclaims the lock', () => {
     expect((mockEngine.speak.mock.calls as unknown as string[][]).map((c) => c[0])).toContain('Fresh.');
   });
 });
+
+// The device bug: pausing a streaming auto-speak fed the paused engine more segments,
+// which timed out and tripped the 2-failure "engine wedged → release" path — unloading
+// the engine so ALL later playback died. "Stop" (stopStreamingSpeechForTurn) must abort
+// cleanly: no wedge, no release, and remaining tokens for the turn must NOT re-engage.
+describe('streaming state machine — user stops mid-stream', () => {
+  it('aborts cleanly without wedging/releasing the engine, and suppresses the rest of the turn', async () => {
+    speakMode = 'hang'; // engine can't complete (as if paused) — the exact device condition
+    feedStreamingText('One. Two. Three. ');
+    await flush();
+    expect(isStreamingSpeechActive()).toBe(true);
+    expect(mockEngine.speak).toHaveBeenCalledTimes(1); // segment 1 in flight (hung)
+
+    stopStreamingSpeechForTurn(); // user hits STOP mid-segment
+
+    // Let the hung speak time out; the orphaned drain must exit, not advance/wedge.
+    await new Promise((r) => setTimeout(r, 80));
+    await flush();
+
+    expect(names()).toContain('stopStreamingSpeechForTurn (user stop — suppress rest of turn)');
+    expect(names()).not.toContain('stream drain ABORT: engine wedged → release for fresh remount');
+    expect(mockEngine.release).not.toHaveBeenCalled();
+    expect(mockEngine.speak).toHaveBeenCalledTimes(1); // never advanced to segment 2
+    expect(isStreamingSpeechActive()).toBe(false);
+
+    // More tokens on the SAME turn must not restart speech.
+    feedStreamingText('One. Two. Three. Four. ');
+    await flush();
+    expect(isStreamingSpeechActive()).toBe(false);
+    expect(mockEngine.speak).toHaveBeenCalledTimes(1);
+  });
+
+  it('a new turn (resetStreamingSpeech) clears the suppression and streams again', async () => {
+    feedStreamingText('Alpha. ');
+    await flush();
+    stopStreamingSpeechForTurn();
+    feedStreamingText('Alpha. Beta. ');
+    await flush();
+    expect(isStreamingSpeechActive()).toBe(false); // still suppressed this turn
+
+    resetStreamingSpeech(); // audio.stop fires this at the next turn
+    feedStreamingText('Gamma. ');
+    await flush();
+    expect(isStreamingSpeechActive()).toBe(true);
+  });
+});

__tests__/integration/generation/imageGenerationFlow.test.ts

@@ -20,6 +20,8 @@ import {
 } from '../../utils/testHelpers';
 import { createONNXImageModel, createGeneratedImage, createMessage } from '../../utils/factories';
 import { Message } from '../../../src/types';
+import { useModelFailureStore } from '../../../src/stores/modelFailureStore';
+import { OverridableMemoryError } from '../../../src/services/modelLoadErrors';
 
 // Mock the services
 jest.mock('../../../src/services/localDreamGenerator');
@@ -456,8 +458,9 @@ describe('Image Generation Flow Integration', () => {
 
       await imageGenerationService.generateImage({ prompt: 'Test' });
 
-      // Should have tried to load model
-      expect(mockActiveModelService.loadImageModel).toHaveBeenCalledWith('img-model-1');
+      // Should have tried to load model (override opts threaded through — undefined on
+      // the normal path, { override: true } only on a Load-Anyway retry).
+      expect(mockActiveModelService.loadImageModel).toHaveBeenCalledWith('img-model-1', undefined, undefined);
     });
 
     it('should reload model if threads changed', async () => {
@@ -1581,6 +1584,154 @@ describe('Image Generation Flow Integration', () => {
       // A successful generation warms the model so the notice never shows again.
       expect(useAppStore.getState().warmedImageModels).toContain(imageModel.id);
     });
+
+    it('once steps advance on a first run, the label says "Generating" — NOT "GPU optimization in progress"', async () => {
+      const imageModel = setupImageModelState();
+      useAppStore.setState({ warmedImageModels: [] }); // first run
+      useAppStore.setState({ settings: { ...useAppStore.getState().settings, imageUseOpenCL: false } });
+      mockLocalDreamService.isModelLoaded.mockResolvedValue(true);
+      mockLocalDreamService.getLoadedThreads.mockReturnValue(4);
+      // Fire progress past the first step so we exercise the mid-generation label.
+      mockLocalDreamService.generateImage.mockImplementation(async (_p, onProgress) => {
+        onProgress?.({ step: 1, totalSteps: 20, progress: 0.05 });
+        onProgress?.({ step: 6, totalSteps: 20, progress: 0.3 });
+        return { id: 'img-1', prompt: 'test', imagePath: '/path/img.png', width: 512, height: 512, steps: 20, seed: 1, modelId: imageModel.id, createdAt: new Date().toISOString() };
+      });
+
+      const statusUpdates: (string | null)[] = [];
+      const unsub = imageGenerationService.subscribe(s => { if (s.status) statusUpdates.push(s.status); });
+      await imageGenerationService.generateImage({ prompt: 'a dog' });
+      unsub();
+
+      // The misleading "GPU optimization in progress (N/steps)" must be gone...
+      expect(statusUpdates.some(s => s?.includes('GPU optimization in progress'))).toBe(false);
+      // ...replaced by an honest "Generating image (6/20)" (with a one-time optimize aside).
+      const midStep = statusUpdates.find(s => s?.includes('6/20'));
+      expect(midStep).toContain('Generating image');
+      expect(midStep).toContain('one-time');
+    });
+  });
+
+  // ============================================================================
+  // Load Anyway override parity: an OVERRIDABLE memory-gate failure on image load
+  // must surface the same "Load Anyway" the text path has, and invoking it must
+  // re-run the load forcing past the budget. Before the fix, imageGenerationService
+  // stringified the typed OverridableMemoryError, so the override was never offered.
+  // ============================================================================
+  describe('size floor (never generate at a garbage sub-256 resolution)', () => {
+    it('floors a stale 128 setting up to 256 before it reaches the native pipeline', async () => {
+      const imageModel = setupImageModelState();
+      // Simulate the on-device state: user had dragged size down to 128 (garbage for SD1.5).
+      useAppStore.setState({ settings: { ...useAppStore.getState().settings, imageWidth: 128, imageHeight: 128 } });
+      mockActiveModelService.getActiveModels.mockReturnValue({
+        text: { model: null, isLoaded: false, isLoading: false },
+        image: { model: imageModel, isLoaded: true, isLoading: false },
+      });
+
+      await imageGenerationService.generateImage({ prompt: 'a dog' });
+
+      expect(mockLocalDreamService.generateImage).toHaveBeenCalledWith(
+        expect.objectContaining({ width: 256, height: 256 }),
+        expect.any(Function),
+        expect.any(Function),
+      );
+    });
+
+    it('passes a valid 512 through unchanged', async () => {
+      const imageModel = setupImageModelState();
+      useAppStore.setState({ settings: { ...useAppStore.getState().settings, imageWidth: 512, imageHeight: 512 } });
+      mockActiveModelService.getActiveModels.mockReturnValue({
+        text: { model: null, isLoaded: false, isLoading: false },
+        image: { model: imageModel, isLoaded: true, isLoading: false },
+      });
+
+      await imageGenerationService.generateImage({ prompt: 'a dog' });
+
+      expect(mockLocalDreamService.generateImage).toHaveBeenCalledWith(
+        expect.objectContaining({ width: 512, height: 512 }),
+        expect.any(Function),
+        expect.any(Function),
+      );
+    });
+  });
+
+  describe('Load Anyway override (overridable memory gate)', () => {
+    beforeEach(() => useModelFailureStore.getState().clear());
+
+    it('offers Load Anyway on an overridable gate, and the retry forces the load with { override: true }', async () => {
+      setupImageModelState();
+      // Model not resident → a load is attempted; the gate rejects it as overridable,
+      // then the forced retry succeeds.
+      mockLocalDreamService.isModelLoaded.mockResolvedValue(false);
+      mockActiveModelService.loadImageModel
+        .mockRejectedValueOnce(
+          new OverridableMemoryError('Not enough memory to load Test Model. Free up space or choose a smaller model.'),
+        )
+        .mockResolvedValue();
+
+      const result = await imageGenerationService.generateImage({ prompt: 'a fox in snow' });
+      expect(result).toBeNull();
+
+      const failure = useModelFailureStore.getState().failures.find(f => f.modelType === 'image');
+      expect(failure).toBeDefined();
+      // The discriminant survived the layers (this is the exact regression).
+      expect(failure!.overridable).toBe(true);
+      expect(typeof failure!.onLoadAnyway).toBe('function');
+
+      // Invoke "Load Anyway" → it must re-attempt the load FORCING past the budget.
+      failure!.onLoadAnyway!();
+      for (let i = 0; i < 6; i++) await flushPromises();
+
+      expect(mockActiveModelService.loadImageModel).toHaveBeenLastCalledWith(
+        'img-model-1',
+        undefined,
+        { override: true },
+      );
+    });
+
+    it('stops offering Load Anyway once the override retry also fails (no repeated no-op)', async () => {
+      setupImageModelState();
+      mockLocalDreamService.isModelLoaded.mockResolvedValue(false);
+      // First load: an overridable gate. The override retry hits the survival floor, so
+      // the service reports it as a NON-overridable hard limit (a plain Error) — the exact
+      // behavior checkImageModelCanLoad now produces under { override: true }.
+      mockActiveModelService.loadImageModel
+        .mockRejectedValueOnce(
+          new OverridableMemoryError('Not enough memory to load Test Model. Free up space or choose a smaller model.'),
+        )
+        .mockRejectedValue(
+          new Error('Not enough memory to load Test Model, even after freeing other models. Close other apps or choose a smaller model.'),
+        );
+
+      await imageGenerationService.generateImage({ prompt: 'a fox' });
+      const first = useModelFailureStore.getState().failures.find(f => f.modelType === 'image');
+      expect(first!.overridable).toBe(true);
+      expect(typeof first!.onLoadAnyway).toBe('function');
+
+      // Press "Load Anyway" → the forced retry fails as a hard limit.
+      first!.onLoadAnyway!();
+      for (let i = 0; i < 6; i++) await flushPromises();
+
+      // The card must NOT keep offering "Load Anyway" — the action would be a no-op.
+      const after = useModelFailureStore.getState().failures.find(f => f.modelType === 'image');
+      expect(after).toBeDefined();
+      expect(after!.overridable).toBeFalsy();
+      expect(after!.onLoadAnyway).toBeUndefined();
+    });
+
+    it('does NOT offer Load Anyway for a NON-overridable load failure (false branch)', async () => {
+      setupImageModelState();
+      mockLocalDreamService.isModelLoaded.mockResolvedValue(false);
+      mockActiveModelService.loadImageModel.mockRejectedValue(new Error('Pipeline failed: model corrupt'));
+
+      const result = await imageGenerationService.generateImage({ prompt: 'a fox' });
+      expect(result).toBeNull();
+
+      const failure = useModelFailureStore.getState().failures.find(f => f.modelType === 'image');
+      expect(failure).toBeDefined();
+      expect(failure!.overridable).toBeFalsy();
+      expect(failure!.onLoadAnyway).toBeUndefined();
+    });
   });
 
   describe('_ensureImageModelLoaded with null activeImageModelId', () => {
@@ -1590,7 +1741,7 @@ describe('Image Generation Flow Integration', () => {
       mockLocalDreamService.getLoadedModelPath.mockResolvedValue(null);
       mockLocalDreamService.getLoadedThreads.mockReturnValue(4);
 
-      const result = await (imageGenerationService as any)._ensureImageModelLoaded(null, fakeModel, 4);
+      const result = await (imageGenerationService as any)._ensureImageModelLoaded(null, fakeModel, { desiredThreads: 4 });
 
       expect(result).toBe(false);
       expect(imageGenerationService.getState().error).toBe('No image model selected');

__tests__/integration/models/activeModelService.test.ts

@@ -27,6 +27,15 @@ import { createDownloadedModel, createONNXImageModel, createDeviceInfo } from '.
 jest.mock('../../../src/services/llm');
 jest.mock('../../../src/services/localDreamGenerator');
 jest.mock('../../../src/services/hardware');
+// Integrity is a boundary for these residency/memory tests (the model files aren't laid
+// down on a real disk here). The completeness rule has its own dedicated tests.
+jest.mock('../../../src/utils/imageModelIntegrity', () => ({
+  validateImageModelDir: jest.fn(async () => ({ complete: true, missing: [] })),
+  ensureImageExtractionComplete: jest.fn(async () => {}),
+}));
+
+import { validateImageModelDir } from '../../../src/utils/imageModelIntegrity';
+import { ImageModelIncompleteError } from '../../../src/services/modelLoadErrors';
 
 const mockLlmService = llmService as jest.Mocked<typeof llmService>;
 const mockLocalDreamService = localDreamGeneratorService as jest.Mocked<typeof localDreamGeneratorService>;
@@ -456,6 +465,19 @@ describe('ActiveModelService Integration', () => {
         { backend: 'auto', cpuOnly: false, attentionVariant: undefined },
       );
     });
+
+    it('refuses an INCOMPLETE model — throws ImageModelIncompleteError and never reaches native load (B respects the verdict)', async () => {
+      const broken = createONNXImageModel({ id: 'img-broken', backend: 'mnn' });
+      useAppStore.setState({ downloadedImageModels: [broken], settings: { imageThreads: 4 } as any });
+      mockLocalDreamService.isModelLoaded.mockResolvedValue(false);
+      // Force the REAL loadImageModel to see an INCOMPLETE verdict from the integrity
+      // boundary. Assert the CONSEQUENCE (throws + native load never happens), not the call —
+      // this catches a caller that queries integrity but ignores `complete: false`.
+      (validateImageModelDir as jest.Mock).mockResolvedValueOnce({ complete: false, missing: ['pos_emb.bin', 'clip_v2.mnn.weight'] });
+
+      await expect(activeModelService.loadImageModel('img-broken')).rejects.toBeInstanceOf(ImageModelIncompleteError);
+      expect(mockLocalDreamService.loadModel).not.toHaveBeenCalled();
+    });
   });
 
   describe('extreme mode (aggressive) — single-model switching text/image/STT', () => {

__tests__/integration/models/imageDownloadRecovery.test.ts

@@ -12,6 +12,13 @@ const mockModelManager = {
   addDownloadedImageModel: jest.fn(),
 };
 
+// Integrity is a boundary for this recovery test (it exercises the resume/move
+// orchestration, not the completeness rule — that has its own dedicated tests).
+jest.mock('../../../src/utils/imageModelIntegrity', () => ({
+  validateImageModelDir: jest.fn(async () => ({ complete: true, missing: [] })),
+  ensureImageExtractionComplete: jest.fn(async () => {}),
+}));
+
 jest.mock('../../../src/services/backgroundDownloadService', () => ({
   backgroundDownloadService: mockBackgroundDownloadService,
 }));