fix: recover SmolLM/Gemma stability + speed up model downloads#7
Conversation
Prevent KV overflow and ChatML regurgitation on SmolLM, honor preferred/custom model selection, gate gated HF downloads, and stop modal/edit dispose crashes after import. Co-authored-by: Cursor <cursoragent@cursor.com>
Per-chunk UI progress callbacks starved the event loop on large HF downloads. Report percent only when it changes and allow 45 minutes for Gemma-scale installs. Co-authored-by: Cursor <cursoragent@cursor.com>
|
Reviewer's GuideRefactors model loading, teardown, and session management to make SmolLM/Gemma inference more stable, gate Gemma downloads behind HF tokens, throttle large model download progress updates, and improve UX flows around custom model import, screenshot questions, and error/empty responses. Sequence diagram for updated model load with timeout and stale completion handlingsequenceDiagram
participant ModelOrchestrator
participant MemoryDiagnosticsService
participant PlatformAdaptationService
participant FlutterGemma
participant InferenceModel
ModelOrchestrator->>MemoryDiagnosticsService: readTotalMemMb()
MemoryDiagnosticsService-->>ModelOrchestrator: lastTotalMemMb
ModelOrchestrator->>PlatformAdaptationService: preferredBackendFor(modelToLoad)
PlatformAdaptationService-->>ModelOrchestrator: PreferredBackend
ModelOrchestrator->>ModelOrchestrator: _loadActiveModelWithTimeout(modelToLoad, preferredBackend, supportImage)
activate ModelOrchestrator
ModelOrchestrator->>ModelOrchestrator: _loadEpoch++ (capture epoch)
ModelOrchestrator->>FlutterGemma: getActiveModel(maxTokens, preferredBackend, supportImage, maxNumImages)
activate FlutterGemma
alt load completes before timeout
FlutterGemma-->>ModelOrchestrator: InferenceModel
deactivate FlutterGemma
ModelOrchestrator->>ModelOrchestrator: check epoch == _loadEpoch
alt epoch unchanged
ModelOrchestrator-->>ModelOrchestrator: return InferenceModel
else epoch changed (stale)
ModelOrchestrator->>InferenceModel: close()
InferenceModel-->>ModelOrchestrator: done
ModelOrchestrator->>ModelOrchestrator: _clearActiveInferenceIdentity()
ModelOrchestrator-->>ModelOrchestrator: throw TimeoutException
end
else TimeoutException
ModelOrchestrator->>ModelOrchestrator: _loadEpoch++
ModelOrchestrator-->>ModelOrchestrator: rethrow TimeoutException
ModelOrchestrator->>FlutterGemma: getActiveModel(...) future.then(...)
activate FlutterGemma
FlutterGemma-->>ModelOrchestrator: late InferenceModel
deactivate FlutterGemma
ModelOrchestrator->>InferenceModel: close()
InferenceModel-->>ModelOrchestrator: done
ModelOrchestrator->>ModelOrchestrator: _clearActiveInferenceIdentity()
end
deactivate ModelOrchestrator
Note over ModelOrchestrator,FlutterGemma: _teardownPartialLoad() also bumps _loadEpoch and closes _activeModel on partial failures
File-Level Changes
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
There was a problem hiding this comment.
Hey - I've left some high level feedback:
- Several
AgentDebugLog.logcalls areawaited directly in hot paths likeprocessMessageand model loading; consider making these fire-and-forget (with appropriate error swallowing) to avoid adding avoidable latency or blocking user-visible flows on logging I/O. - The HuggingFace gating logic is split between
requiresHuggingFaceAuth(NovaModel)andurlRequiresHuggingFaceAuth(String)and used in different places; consolidating these into a single, model-centric check would reduce duplication and the risk of the two paths diverging over time. - The
_loadEpochmechanism is updated in multiple code paths (_loadActiveModelWithTimeout,_teardownPartialLoad, timeout handlers); it may be worth adding a brief comment or small helper around epoch increments to make the intended sequencing/race guarantees clearer for future maintainers.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- Several `AgentDebugLog.log` calls are `await`ed directly in hot paths like `processMessage` and model loading; consider making these fire-and-forget (with appropriate error swallowing) to avoid adding avoidable latency or blocking user-visible flows on logging I/O.
- The HuggingFace gating logic is split between `requiresHuggingFaceAuth(NovaModel)` and `urlRequiresHuggingFaceAuth(String)` and used in different places; consolidating these into a single, model-centric check would reduce duplication and the risk of the two paths diverging over time.
- The `_loadEpoch` mechanism is updated in multiple code paths (`_loadActiveModelWithTimeout`, `_teardownPartialLoad`, timeout handlers); it may be worth adding a brief comment or small helper around epoch increments to make the intended sequencing/race guarantees clearer for future maintainers.Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces several enhancements and stability fixes for model management and inference orchestration, including session invalidation for replay on message edits, HuggingFace token checks for gated Gemma models, optimized download progress reporting, and forcing the CPU backend for SmolLM on low-RAM Android devices. Feedback on these changes highlights critical asynchronous issues, such as a race condition in _teardownActiveModel where the active model is not cleared synchronously before closing, a potential RangeError when mutating the message list after an asynchronous gap, a missing mounted check before calling setState in the settings screen, and a bug where late timeout cleanups could inadvertently clear the active identity of a newly loaded model.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| Future<void> _teardownActiveModel({NovaModel? keepIfSameType}) async { | ||
| _activeChat = null; | ||
| if (_activeModel == null) return; | ||
| if (_activeModel == null) { | ||
| await _clearActiveInferenceIdentity(); | ||
|
|
||
| return; | ||
| } | ||
| final shouldClose = | ||
| keepIfSameType == null || | ||
| _activeModelType == null || | ||
| keepIfSameType != _activeModelType; | ||
| if (!shouldClose) return; | ||
| _activeModel!.close().catchError((_) {}); | ||
| try { | ||
| await _activeModel!.close(); | ||
| } catch (_) {} | ||
| _activeModel = null; | ||
| _activeModelType = null; | ||
| _activeModelSupportsImage = false; | ||
| _clearActiveInferenceIdentity(); | ||
| await _clearActiveInferenceIdentity(); | ||
| } |
There was a problem hiding this comment.
There is a critical race condition in _teardownActiveModel. Because this method is asynchronous and unawaited in multiple places (e.g., lines 330, 349, 756), _activeModel remains non-null while await _activeModel!.close() is executing. If another operation (such as _getOrCreateModel or another teardown) runs concurrently, it will see _activeModel as non-null and potentially attempt to call close() on it again. Calling close() concurrently on the same native model instance can cause severe native crashes (SIGABRT/double-free).
To fix this, clear _activeModel, _activeModelType, and _activeModelSupportsImage synchronously before awaiting the asynchronous close() call.
Future<void> _teardownActiveModel({NovaModel? keepIfSameType}) async {
_activeChat = null;
if (_activeModel == null) {
await _clearActiveInferenceIdentity();
return;
}
final shouldClose =
keepIfSameType == null ||
_activeModelType == null ||
keepIfSameType != _activeModelType;
if (!shouldClose) return;
final modelToClose = _activeModel;
_activeModel = null;
_activeModelType = null;
_activeModelSupportsImage = false;
try {
await modelToClose!.close();
} catch (_) {}
await _clearActiveInferenceIdentity();
}| await TtsService.instance.stop(); | ||
| if (!mounted) return; | ||
| setState(() { | ||
| _messages.removeRange(userIndex, _messages.length); | ||
| }); |
There was a problem hiding this comment.
There is an asynchronous gap (await TtsService.instance.stop()) before _messages.removeRange is called. During this await, the user could have cleared the chat history or switched sessions, meaning _messages could be empty or smaller than userIndex. Calling removeRange without validating userIndex against the current _messages.length will throw a RangeError and crash the UI.
Validate that userIndex is still within bounds of _messages after the asynchronous gap.
| await TtsService.instance.stop(); | |
| if (!mounted) return; | |
| setState(() { | |
| _messages.removeRange(userIndex, _messages.length); | |
| }); | |
| await TtsService.instance.stop(); | |
| if (!mounted) return; | |
| if (userIndex < 0 || userIndex >= _messages.length) return; | |
| setState(() { | |
| _messages.removeRange(userIndex, _messages.length); | |
| }); |
| if (ModelHuggingFaceURLs.requiresHuggingFaceAuth(model) && | ||
| !await ModelManager.hasHuggingFaceToken()) { | ||
| if (!context.mounted) return; | ||
| ScaffoldMessenger.of(context).showSnackBar( | ||
| SnackBar( | ||
| content: const Text( | ||
| 'Gemma models need a HuggingFace token. Add one in Settings.', | ||
| ), | ||
| backgroundColor: Colors.orange[800], | ||
| action: SnackBarAction( | ||
| label: 'Add token', | ||
| textColor: Colors.white, | ||
| onPressed: () => _showHfTokenDialog(context), | ||
| ), | ||
| ), | ||
| ); | ||
|
|
||
| return; | ||
| } | ||
|
|
||
| setState(() => _installStatus = 'Downloading ${model.displayName}...'); |
There was a problem hiding this comment.
In _downloadModel, ModelManager.hasHuggingFaceToken() is awaited. If the token check succeeds, the code proceeds to call setState at line 1289 without checking if the widget is still mounted. If the user exits the settings screen while the token check is pending, this will trigger a "setState() called after dispose()" crash.
Ensure mounted is checked before calling setState after the asynchronous gap.
if (ModelHuggingFaceURLs.requiresHuggingFaceAuth(model) &&
!await ModelManager.hasHuggingFaceToken()) {
if (!context.mounted) return;
ScaffoldMessenger.of(context).showSnackBar(
SnackBar(
content: const Text(
'Gemma models need a HuggingFace token. Add one in Settings.',
),
backgroundColor: Colors.orange[800],
action: SnackBarAction(
label: 'Add token',
textColor: Colors.white,
onPressed: () => _showHfTokenDialog(context),
),
),
);
return;
}
if (!mounted) return;
setState(() => _installStatus = 'Downloading ${model.displayName}...');| } on TimeoutException { | ||
| _loadEpoch++; | ||
| unawaited( | ||
| future | ||
| .then((lateModel) async { | ||
| try { | ||
| await lateModel.close(); | ||
| } catch (_) {} | ||
| await _clearActiveInferenceIdentity(); | ||
| }) | ||
| .catchError((_) {}), | ||
| ); | ||
| rethrow; | ||
| } |
There was a problem hiding this comment.
In _loadActiveModelWithTimeout, if a model load times out, a late completion callback is registered to close the late model and call _clearActiveInferenceIdentity(). However, if a new model has already successfully loaded in the meantime, calling _clearActiveInferenceIdentity() will clear the active identity of the newly loaded model, breaking its session.
Only call _clearActiveInferenceIdentity() in the late completion callback if no other model is currently active (i.e., _activeModel == null).
} on TimeoutException {
_loadEpoch++;
unawaited(
future
.then((lateModel) async {
try {
await lateModel.close();
} catch (_) {}
if (_activeModel == null) {
await _clearActiveInferenceIdentity();
}
})
.catchError((_) {}),
);
rethrow;
}
Code Review SummaryStatus: 4 Issues Found | Recommendation: Address before merge Overview
Issue Details (click to expand)CRITICAL
HIGH
MEDIUM
Files Reviewed (15 files)
Fix these issues in Kilo Cloud Reviewed by step-3.7-flash · Input: 294.5K · Output: 59.3K · Cached: 7.9M |
1 participant
Summary
Test plan
Made with Cursor
Summary by Sourcery
Stabilize SmolLM and Gemma sessions and improve reliability of large model installs and chats.
Bug Fixes:
Enhancements:
Documentation:
Tests: