fix(ai): harden model switching, downloads, and migration

- cancel and await in-flight model downloads before switching configs\n- ignore stale download callbacks so old transfers cannot mutate the new model state\n- dispose LlamaEngine instances on unload and error paths in both chat and embedding services\n- reject non-success HTTP responses before caching model files and normalize forced-cancel errors\n- share GGUF URL validation and filename derivation between settings and model config creation\n- move embedding migration into a reusable post-load flow used by splash and settings\n- keep flutter analyze and both lite/full release APK builds passing

Author: priyavratuniyal

lib/main.dart

@@ -22,6 +22,7 @@ late BiomarkerDictionary biomarkerDictionary;
 late LlmService llmService;
 late LlmEmbeddingService embeddingService;
 late VectorStoreService vectorStoreService;
+Future<void>? _embeddingMigrationTask;
 
 /// Whether AI features are enabled. Set by the entry point
 /// (main_full.dart vs main_lite.dart).
@@ -38,6 +39,38 @@ Future<void> appMain({required bool aiEnabled}) async {
   runApp(const KoshikaApp());
 }
 
+Future<void> migrateEmbeddingsIfNeeded() {
+  if (!kAiEnabled || !embeddingService.isLoaded) {
+    return Future.value();
+  }
+
+  return _embeddingMigrationTask ??= _runEmbeddingMigration().whenComplete(() {
+    _embeddingMigrationTask = null;
+  });
+}
+
+Future<void> _runEmbeddingMigration() async {
+  try {
+    final allResults = objectbox.biomarkerResultBox.getAll();
+    if (allResults.isEmpty) return;
+
+    final needsMigration = allResults.any(
+      (result) =>
+          result.embedding == null ||
+          result.embedding!.isEmpty ||
+          result.embedding!.length != 384,
+    );
+    if (!needsMigration) return;
+
+    debugPrint(
+      'KoshikaApp: migrating ${allResults.length} embeddings to 384-dim',
+    );
+    await vectorStoreService.rebuildIndex(allResults);
+  } catch (e) {
+    debugPrint('Embedding migration failed (non-fatal): $e');
+  }
+}
+
 class KoshikaApp extends StatelessWidget {
   const KoshikaApp({super.key});
 

lib/models/llm_model_config.dart

@@ -1,3 +1,59 @@
+class CustomModelUrlDetails {
+  final String normalizedUrl;
+  final String filename;
+
+  const CustomModelUrlDetails({
+    required this.normalizedUrl,
+    required this.filename,
+  });
+
+  String get suggestedName =>
+      filename.replaceFirst(RegExp(r'\.gguf$', caseSensitive: false), '');
+}
+
+String? _extractGgufFilename(Uri uri) {
+  if (uri.pathSegments.isNotEmpty) {
+    final candidate = uri.pathSegments.last;
+    if (candidate.isNotEmpty && candidate.toLowerCase().endsWith('.gguf')) {
+      return candidate;
+    }
+  }
+
+  for (final values in uri.queryParametersAll.values) {
+    for (final value in values) {
+      final candidate = _extractGgufFilenameFromValue(value);
+      if (candidate != null) return candidate;
+    }
+  }
+
+  return null;
+}
+
+String? _extractGgufFilenameFromValue(String value) {
+  final trimmed = value.trim();
+  if (trimmed.isEmpty) return null;
+
+  final uri = Uri.tryParse(trimmed);
+  final candidate = uri != null && uri.pathSegments.isNotEmpty
+      ? uri.pathSegments.last
+      : trimmed.split('/').last;
+
+  if (candidate.isEmpty || !candidate.toLowerCase().endsWith('.gguf')) {
+    return null;
+  }
+
+  return candidate;
+}
+
+String _filenameFromDownloadUrl(
+  String downloadUrl, {
+  required String fallbackId,
+}) {
+  final uri = Uri.tryParse(downloadUrl);
+  if (uri == null) return '$fallbackId.gguf';
+  return _extractGgufFilename(uri) ?? '$fallbackId.gguf';
+}
+
 /// Configuration for a downloadable GGUF model (curated or custom).
 class LlmModelConfig {
   final String id;
@@ -18,8 +74,7 @@ class LlmModelConfig {
 
   /// Derive the on-disk filename from the download URL.
   String get filename {
-    final uri = Uri.parse(downloadUrl);
-    return uri.pathSegments.isNotEmpty ? uri.pathSegments.last : '$id.gguf';
+    return _filenameFromDownloadUrl(downloadUrl, fallbackId: id);
   }
 
   /// Human-readable size string.
@@ -86,16 +141,43 @@ abstract final class LlmModelRegistry {
     required String name,
     required String downloadUrl,
   }) {
+    final parsed = inspectCustomDownloadUrl(downloadUrl);
+    final displayName = name.trim().isNotEmpty
+        ? name.trim()
+        : parsed.suggestedName;
+
     return LlmModelConfig(
       id: 'custom',
-      name: name,
-      downloadUrl: downloadUrl,
+      name: displayName,
+      downloadUrl: parsed.normalizedUrl,
       estimatedSizeMB: 0,
       description: 'Custom GGUF model',
       isCustom: true,
     );
   }
 
+  static CustomModelUrlDetails inspectCustomDownloadUrl(String rawUrl) {
+    final normalizedUrl = rawUrl.trim();
+    if (normalizedUrl.isEmpty) {
+      throw ArgumentError('URL is required');
+    }
+
+    final uri = Uri.tryParse(normalizedUrl);
+    if (uri == null || uri.scheme != 'https' || !uri.hasAuthority) {
+      throw ArgumentError('URL must be a valid HTTPS link');
+    }
+
+    final filename = _extractGgufFilename(uri);
+    if (filename == null) {
+      throw ArgumentError('URL must point to a .gguf file');
+    }
+
+    return CustomModelUrlDetails(
+      normalizedUrl: normalizedUrl,
+      filename: filename,
+    );
+  }
+
   /// Look up a curated model by id. Returns null for unknown ids.
   static LlmModelConfig? findById(String id) {
     for (final m in curated) {

lib/screens/settings_screen.dart

@@ -117,6 +117,11 @@ class _SettingsScreenState extends State<SettingsScreen> {
     if (result != null) await _onModelSelected(result);
   }
 
+  Future<void> _loadEmbeddingModel() async {
+    await embeddingService.loadModel();
+    await migrateEmbeddingsIfNeeded();
+  }
+
   // ─── Data Actions ───────────────────────────────────────────────────
 
   Future<void> _deleteAllData() async {
@@ -252,7 +257,7 @@ class _SettingsScreenState extends State<SettingsScreen> {
             _EmbeddingModelTile(
               modelInfo: _embeddingInfo,
               onDownload: () => embeddingService.downloadModel(),
-              onLoad: () => embeddingService.loadModel(),
+              onLoad: _loadEmbeddingModel,
               onUnload: () => embeddingService.unloadModel(),
             ),
             const SizedBox(height: KoshikaSpacing.sm),
@@ -955,24 +960,14 @@ class _CustomModelDialogState extends State<_CustomModelDialog> {
       setState(() => _urlError = 'URL is required');
       return;
     }
-    if (!url.toLowerCase().endsWith('.gguf')) {
-      setState(() => _urlError = 'URL must point to a .gguf file');
-      return;
-    }
-
-    final displayName = name.isNotEmpty ? name : _filenameFromUrl(url);
-    Navigator.of(
-      context,
-    ).pop(LlmModelRegistry.custom(name: displayName, downloadUrl: url));
-  }
-
-  String _filenameFromUrl(String url) {
     try {
-      final uri = Uri.parse(url);
-      final filename = uri.pathSegments.last;
-      return filename.replaceAll('.gguf', '');
-    } catch (_) {
-      return 'Custom Model';
+      final parsed = LlmModelRegistry.inspectCustomDownloadUrl(url);
+      final displayName = name.isNotEmpty ? name : parsed.suggestedName;
+      Navigator.of(
+        context,
+      ).pop(LlmModelRegistry.custom(name: displayName, downloadUrl: url));
+    } on ArgumentError catch (e) {
+      setState(() => _urlError = e.message?.toString() ?? 'Invalid GGUF URL');
     }
   }
 

lib/screens/splash_screen.dart

@@ -88,8 +88,8 @@ class _SplashScreenState extends State<SplashScreen>
             ModelStatus.ready) {
           // ignore: unawaited_futures
           app_main.embeddingService.loadModel().then((_) {
-            // Re-embed any results that have stale or missing embeddings
-            _migrateEmbeddingsIfNeeded();
+            // Re-embed any results that have stale or missing embeddings.
+            app_main.migrateEmbeddingsIfNeeded();
           });
         }
       }
@@ -120,35 +120,6 @@ class _SplashScreenState extends State<SplashScreen>
     }
   }
 
-  /// One-time migration: recompute all embeddings when the embedding
-  /// dimension changes (768 → 384 after flutter_gemma → llama_cpp_dart).
-  Future<void> _migrateEmbeddingsIfNeeded() async {
-    if (!app_main.kAiEnabled) return;
-    if (!app_main.embeddingService.isLoaded) return;
-
-    try {
-      final allResults = app_main.objectbox.biomarkerResultBox.getAll();
-      if (allResults.isEmpty) return;
-
-      // Check if any result has stale (768-dim) or missing embeddings
-      final needsMigration = allResults.any(
-        (r) =>
-            r.embedding == null ||
-            r.embedding!.isEmpty ||
-            r.embedding!.length != 384,
-      );
-
-      if (needsMigration) {
-        debugPrint(
-          'SplashScreen: migrating ${allResults.length} embeddings to 384-dim',
-        );
-        await app_main.vectorStoreService.rebuildIndex(allResults);
-      }
-    } catch (e) {
-      debugPrint('Embedding migration failed (non-fatal): $e');
-    }
-  }
-
   @override
   void dispose() {
     _controller.dispose();

lib/services/llm_embedding_service.dart

@@ -149,8 +149,7 @@ class LlmEmbeddingService {
       await _engine!.loadModel(path);
       _updateStatus(_modelInfo.copyWith(status: ModelStatus.loaded));
     } catch (e) {
-      await _engine?.dispose();
-      _engine = null;
+      await _disposeEngine();
       _updateStatus(
         _modelInfo.copyWith(
           status: ModelStatus.error,
@@ -185,9 +184,7 @@ class LlmEmbeddingService {
   // ═══════════════════════════════════════════════════════════════════════
 
   Future<void> unloadModel() async {
-    if (_engine != null) {
-      await _engine!.unloadModel();
-    }
+    await _disposeEngine(unloadFirst: true);
     if (_modelInfo.status == ModelStatus.loaded ||
         _modelInfo.status == ModelStatus.loading) {
       _updateStatus(_modelInfo.copyWith(status: ModelStatus.ready));
@@ -203,9 +200,22 @@ class LlmEmbeddingService {
     if (!_statusController.isClosed) _statusController.add(info);
   }
 
-  Future<void> dispose() async {
-    await _engine?.dispose();
+  Future<void> _disposeEngine({bool unloadFirst = false}) async {
+    final engine = _engine;
     _engine = null;
+    if (engine == null) return;
+
+    try {
+      if (unloadFirst && engine.isReady) {
+        await engine.unloadModel();
+      }
+    } finally {
+      await engine.dispose();
+    }
+  }
+
+  Future<void> dispose() async {
+    await _disposeEngine();
     _statusController.close();
   }
 }