Skip to content

Commit bc41e84

Browse files
authored
ADFA-1798 | Add proactive memory warning for custom model loading (#63)
feat(ai-assistant): warn before loading an oversized local model Selecting a .gguf now runs a pre-flight check before the path is persisted, which is what makes ai-core load it. The model's size and its declared shape are weighed against free RAM; if it looks too large the user is asked whether to go ahead, with Cancel as the default and the figures behind the judgement on screen. Declining leaves the previously selected model untouched. The check fails open throughout: an unreadable size, header or memory reading means no warning rather than a wrong one. - ModelMemoryEstimator splits mmap'd weights from the KV cache and compute buffers, sized from the model's own shape when the header provides it. - ModelMemoryGate decides Safe / Risky(TIGHT|INSUFFICIENT) / Unknown. - GgufHeaderReader reads only the metadata block, bounded by a byte budget, an entry cap and an array cap, and attributes shape keys to their own architecture so a multimodal file's clip.* values stay out. - UserConfirmation lets the load coroutine await the user's answer, so the flow stays one readable sequence across a rotation. ADFA-1798
1 parent 06eed4a commit bc41e84

22 files changed

Lines changed: 2342 additions & 43 deletions

File tree

‎ai-assistant/src/main/assets/docs/index.html‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,28 @@ <h2>Choosing a backend</h2>
5656
agent reads are sent to Google over HTTPS.</li>
5757
</ul>
5858

59+
<h2>Will this model fit in memory?</h2>
60+
<p>When you pick a local <code>.gguf</code> file, it is measured against the RAM
61+
free on the device at that moment, before anything is loaded. If it looks too
62+
large you get a warning with the actual figures and two choices:</p>
63+
<ul>
64+
<li><b>Cancel</b> — nothing is saved and nothing is loaded. Whatever model you
65+
had selected before stays in use. Dismissing the warning with Back does the
66+
same thing.</li>
67+
<li><b>Proceed anyway</b> — the model is accepted despite the shortfall. Use
68+
this when you are about to close other apps, for example.</li>
69+
</ul>
70+
<p>The warning quotes two numbers. Memory <b>to load</b> is the weights: these are
71+
memory-mapped, so they need not all fit at once — when they don't, the device pages
72+
them in and out, which is why an oversized model can stall for minutes instead of
73+
failing immediately. Memory <b>to run</b> is the KV cache and compute buffers, which
74+
are ordinary allocations and do have to fit. That is why a model can be reported as
75+
risky rather than impossible: the outcome genuinely depends on how much paging your
76+
device will tolerate.</p>
77+
<p>To fit a large model, close other apps and select it again, or pick a smaller or
78+
more heavily quantized build. A Q4_K_M quantization of a 1–3B model is the most
79+
likely to run comfortably.</p>
80+
5981
<h2>What the agent can do</h2>
6082
<ul>
6183
<li>Read and search files <b>within the current project</b>.</li>

‎ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/AiAssistantPlugin.kt‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,10 @@ class AiAssistantPlugin : IPlugin, UIExtension, DocumentationExtension, Settings
6666
const val TOOLTIP_TAG_SETTINGS_GEMINI_MODEL = "ai_settings_gemini_model"
6767
const val TOOLTIP_TAG_SETTINGS_GET_KEY = "ai_settings_get_free_key"
6868

69+
// Tags for the memory pre-flight warning (see MemoryWarningDialogFragment).
70+
const val TOOLTIP_TAG_MEMORY_PROCEED = "agent_memory_warning_proceed"
71+
const val TOOLTIP_TAG_MEMORY_CANCEL = "agent_memory_warning_cancel"
72+
6973
@Volatile
7074
private var pluginContext: PluginContext? = null
7175

@@ -457,6 +461,39 @@ class AiAssistantPlugin : IPlugin, UIExtension, DocumentationExtension, Settings
457461
models can't generate replies. Larger models are slower and use
458462
more memory; the file is copied into the app's private storage on
459463
first use.</p>
464+
<p>The model is measured against this device's free memory before it
465+
is accepted. If it looks too large you get the figures and a choice
466+
to cancel or continue.</p>
467+
""".trimIndent(),
468+
buttons = listOf(
469+
PluginTooltipButton(description = "AI Assistant guide", uri = "index.html", order = 0)
470+
)
471+
),
472+
PluginTooltipEntry(
473+
tag = TOOLTIP_TAG_MEMORY_PROCEED,
474+
summary = "Load this model anyway, accepting that it may fail or slow the device.",
475+
detail = """
476+
<p>The model's weights plus its working memory look larger than the
477+
RAM free right now. Weights are memory-mapped, so a load can still
478+
succeed by paging — which is why the outcome is a risk rather than a
479+
certainty: it may work, fail quickly, or stall for minutes first.</p>
480+
<p>Use this when you know the numbers are wrong for your situation,
481+
for example because you are about to close other apps.</p>
482+
""".trimIndent(),
483+
buttons = listOf(
484+
PluginTooltipButton(description = "AI Assistant guide", uri = "index.html", order = 0)
485+
)
486+
),
487+
PluginTooltipEntry(
488+
tag = TOOLTIP_TAG_MEMORY_CANCEL,
489+
summary = "Abandon this model; the previously selected one is left untouched.",
490+
detail = """
491+
<p>Nothing is saved and nothing is loaded, so the model you had
492+
selected before stays in use. This is the safe choice, and also what
493+
happens if you dismiss the warning with Back.</p>
494+
<p>To fit a large model, close other apps and pick it again, or
495+
choose a smaller or more heavily quantized build — a Q4_K_M
496+
quantization of a 1–3B model is the most likely to run.</p>
460497
""".trimIndent(),
461498
buttons = listOf(
462499
PluginTooltipButton(description = "AI Assistant guide", uri = "index.html", order = 0)

‎ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/fragments/AiSettingsFragment.kt‎

Lines changed: 55 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,10 @@ import android.widget.*
1717
import androidx.activity.result.contract.ActivityResultContracts
1818
import androidx.annotation.DrawableRes
1919
import androidx.fragment.app.Fragment
20+
import androidx.lifecycle.Lifecycle
2021
import androidx.lifecycle.ViewModelProvider
2122
import androidx.lifecycle.lifecycleScope
23+
import androidx.lifecycle.repeatOnLifecycle
2224
import com.google.android.material.dialog.MaterialAlertDialogBuilder
2325
import com.itsaky.androidide.plugins.PluginContext
2426
import com.itsaky.androidide.plugins.aiassistant.AiAssistantPlugin
@@ -31,6 +33,7 @@ import com.itsaky.androidide.plugins.aiassistant.viewmodel.AiBackend
3133
import com.itsaky.androidide.plugins.aiassistant.viewmodel.AiSettingsViewModel
3234
import com.itsaky.androidide.plugins.aiassistant.viewmodel.EngineState
3335
import com.itsaky.androidide.plugins.aiassistant.viewmodel.ModelLoadingState
36+
import com.itsaky.androidide.plugins.aiassistant.viewmodel.ModelMemoryWarning
3437
import kotlinx.coroutines.launch
3538
import java.text.SimpleDateFormat
3639
import java.util.Date
@@ -42,7 +45,7 @@ import kotlin.math.roundToInt
4245
* chat's own shortcuts. The host mounts it full-screen in PluginScreenActivity, which provides no
4346
* toolbar, so this fragment brings its own app bar and closes by finishing that activity.
4447
*/
45-
class AiSettingsFragment : Fragment() {
48+
class AiSettingsFragment : Fragment(), MemoryWarningDialogFragment.Host {
4649

4750
private lateinit var viewModel: AiSettingsViewModel
4851
private lateinit var settingsToolbar: LinearLayout
@@ -130,6 +133,57 @@ class AiSettingsFragment : Fragment() {
130133
initializeViews(view)
131134
setupToolbar()
132135
setupBackendSelector()
136+
observeMemoryWarnings()
137+
}
138+
139+
/**
140+
* Puts a "this model may not fit" question to the user. Collected under STARTED so the dialog is
141+
* never shown to a stopped fragment; the event waits in the ViewModel until then.
142+
*/
143+
private fun observeMemoryWarnings() {
144+
dropStaleMemoryWarning()
145+
viewLifecycleOwner.lifecycleScope.launch {
146+
viewLifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
147+
viewModel.modelMemoryWarnings.collect(::showMemoryWarning)
148+
}
149+
}
150+
}
151+
152+
/**
153+
* Dismiss a warning dialog the framework restored around a question that no longer exists.
154+
* After process death the load that raised it is gone, so every button on it would be a silent
155+
* no-op — better to take it away than to leave the user pressing a dialog that decides nothing.
156+
*/
157+
private fun dropStaleMemoryWarning() {
158+
if (viewModel.hasPendingMemoryWarning) return
159+
val restored = childFragmentManager.findFragmentByTag(MemoryWarningDialogFragment.TAG)
160+
(restored as? MemoryWarningDialogFragment)?.dismissAllowingStateLoss()
161+
}
162+
163+
/**
164+
* Shown as a child fragment, so it survives rotation and can still reach this host. Must stay
165+
* idempotent: an unanswered question is re-published to every new collector by
166+
* [com.itsaky.androidide.plugins.aiassistant.viewmodel.UserConfirmation].
167+
*
168+
* @param warning the model and the figures to put to the user
169+
*/
170+
private fun showMemoryWarning(warning: ModelMemoryWarning) {
171+
if (childFragmentManager.findFragmentByTag(MemoryWarningDialogFragment.TAG) != null) return
172+
MemoryWarningDialogFragment.newInstance(warning)
173+
.show(childFragmentManager, MemoryWarningDialogFragment.TAG)
174+
}
175+
176+
override fun onModelMemoryDecision(proceed: Boolean) {
177+
viewModel.onMemoryWarningDecision(proceed)
178+
// Not requireContext(): onCancel can reach us as the fragment is going away.
179+
val ctx = context ?: return
180+
if (!proceed) {
181+
Toast.makeText(
182+
ctx,
183+
getString(R.string.llm_memory_warning_declined),
184+
Toast.LENGTH_LONG,
185+
).show()
186+
}
133187
}
134188

135189
override fun onResume() {
Lines changed: 153 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,153 @@
1+
package com.itsaky.androidide.plugins.aiassistant.fragments
2+
3+
import android.app.Dialog
4+
import android.content.DialogInterface
5+
import android.os.Bundle
6+
import android.view.View
7+
import androidx.fragment.app.DialogFragment
8+
import com.google.android.material.dialog.MaterialAlertDialogBuilder
9+
import com.itsaky.androidide.plugins.aiassistant.AiAssistantPlugin
10+
import com.itsaky.androidide.plugins.aiassistant.R
11+
import com.itsaky.androidide.plugins.aiassistant.memory.ModelMemoryGate
12+
import com.itsaky.androidide.plugins.aiassistant.util.ByteSize
13+
import com.itsaky.androidide.plugins.aiassistant.viewmodel.ModelMemoryWarning
14+
import com.itsaky.androidide.plugins.base.PluginFragmentHelper
15+
import com.itsaky.androidide.plugins.services.IdeTooltipService
16+
17+
/**
18+
* Warns that a selected model may be too large for this device, offering **Cancel** — the default
19+
* in every sense visible to the user — or **Proceed anyway** (ADFA-1798). Decisions go to [Host],
20+
* resolved per call so they still arrive after the framework recreates this on a rotation.
21+
*/
22+
class MemoryWarningDialogFragment : DialogFragment() {
23+
24+
private val tooltipService: IdeTooltipService? by lazy {
25+
try {
26+
PluginFragmentHelper.getServiceRegistry(AiAssistantPlugin.PLUGIN_ID)
27+
?.get(IdeTooltipService::class.java)
28+
} catch (e: Exception) {
29+
// Tooltip help is optional; long-press simply shows nothing when it's unavailable.
30+
AiAssistantPlugin.getContext()?.logger
31+
?.warn("MemoryWarningDialogFragment: tooltip service unavailable", e)
32+
null
33+
}
34+
}
35+
36+
/**
37+
* Receives this dialog's outcome. Implemented by the fragment that shows the dialog, which must
38+
* be its **parent** fragment (show it with `childFragmentManager`).
39+
*/
40+
interface Host {
41+
/** @param proceed true to load the model anyway, false to abandon the selection. */
42+
fun onModelMemoryDecision(proceed: Boolean)
43+
}
44+
45+
companion object {
46+
const val TAG = "ModelMemoryWarning"
47+
48+
private const val ARG_MODEL_NAME = "model_name"
49+
private const val ARG_LOAD_BYTES = "load_bytes"
50+
private const val ARG_RUN_BYTES = "run_bytes"
51+
private const val ARG_AVAILABLE_BYTES = "available_bytes"
52+
private const val ARG_SEVERITY = "severity"
53+
54+
/**
55+
* Builds the dialog. Everything it needs is in [getArguments], so the framework can recreate
56+
* it after a configuration change while the answer channel stays on the ViewModel.
57+
*
58+
* @param warning the model and the figures to show
59+
*/
60+
fun newInstance(warning: ModelMemoryWarning): MemoryWarningDialogFragment =
61+
MemoryWarningDialogFragment().apply {
62+
arguments = Bundle().apply {
63+
putString(ARG_MODEL_NAME, warning.modelName)
64+
putLong(ARG_LOAD_BYTES, warning.loadBytes)
65+
putLong(ARG_RUN_BYTES, warning.runBytes)
66+
putLong(ARG_AVAILABLE_BYTES, warning.availableBytes)
67+
putString(ARG_SEVERITY, warning.severity.name)
68+
}
69+
}
70+
}
71+
72+
override fun onCreateDialog(savedInstanceState: Bundle?): Dialog {
73+
val args = arguments ?: Bundle.EMPTY
74+
val messageId = when (severity()) {
75+
ModelMemoryGate.Severity.INSUFFICIENT -> R.string.llm_memory_warning_insufficient
76+
ModelMemoryGate.Severity.TIGHT -> R.string.llm_memory_warning_tight
77+
}
78+
val message = getString(
79+
messageId,
80+
args.getString(ARG_MODEL_NAME).orEmpty(),
81+
ByteSize.format(args.getLong(ARG_LOAD_BYTES)),
82+
ByteSize.format(args.getLong(ARG_RUN_BYTES)),
83+
ByteSize.format(args.getLong(ARG_AVAILABLE_BYTES)),
84+
)
85+
86+
val dialog = MaterialAlertDialogBuilder(requireContext())
87+
.setTitle(getString(R.string.llm_memory_warning_title))
88+
.setMessage(message)
89+
// Cancel takes the positive slot: Material emphasizes that button.
90+
.setPositiveButton(getString(R.string.llm_memory_warning_cancel)) { _, _ -> decide(false) }
91+
.setNegativeButton(getString(R.string.llm_memory_warning_proceed)) { _, _ -> decide(true) }
92+
.create()
93+
94+
// Bound after show(): the buttons don't exist before it.
95+
dialog.setOnShowListener {
96+
val cancelButton = dialog.getButton(Dialog.BUTTON_POSITIVE)
97+
wireTooltip(cancelButton, AiAssistantPlugin.TOOLTIP_TAG_MEMORY_CANCEL)
98+
wireTooltip(
99+
dialog.getButton(Dialog.BUTTON_NEGATIVE),
100+
AiAssistantPlugin.TOOLTIP_TAG_MEMORY_PROCEED,
101+
)
102+
// So Enter, D-pad and a screen reader all reach the safe choice first.
103+
cancelButton?.requestFocus()
104+
}
105+
106+
return dialog
107+
}
108+
109+
/**
110+
* Back press and outside tap both mean "don't risk it". Overridden rather than set on the
111+
* builder: [DialogFragment] replaces a builder cancel listener in `prepareDialog`, so the
112+
* awaiting caller would never learn the model was declined.
113+
*
114+
* @param dialog the dialog that was cancelled
115+
*/
116+
override fun onCancel(dialog: DialogInterface) {
117+
super.onCancel(dialog)
118+
decide(false)
119+
}
120+
121+
/** Falls back to the milder wording, so a bad argument can't overstate the risk. */
122+
private fun severity(): ModelMemoryGate.Severity = try {
123+
ModelMemoryGate.Severity.valueOf(
124+
arguments?.getString(ARG_SEVERITY) ?: ModelMemoryGate.Severity.TIGHT.name
125+
)
126+
} catch (e: IllegalArgumentException) {
127+
ModelMemoryGate.Severity.TIGHT
128+
}
129+
130+
/**
131+
* Delivers the outcome to the host fragment. Resolved on each call rather than captured at
132+
* construction, so it still works on the instance the framework recreated after a rotation.
133+
*/
134+
private fun decide(proceed: Boolean) {
135+
val host = parentFragment as? Host
136+
if (host == null) {
137+
// Nothing else reports this, and the load that raised the warning waits forever on it.
138+
AiAssistantPlugin.getContext()?.logger
139+
?.warn("MemoryWarningDialogFragment: no Host parent; decision dropped")
140+
return
141+
}
142+
host.onModelMemoryDecision(proceed)
143+
}
144+
145+
/** Shows this plugin's tooltip for [tag] when [view] is long-pressed (Tier 1/2 + guide). */
146+
private fun wireTooltip(view: View?, tag: String) {
147+
view?.setOnLongClickListener { anchor ->
148+
val service = tooltipService ?: return@setOnLongClickListener false
149+
service.showTooltip(anchor, AiAssistantPlugin.TOOLTIP_CATEGORY, tag)
150+
true
151+
}
152+
}
153+
}
Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
package com.itsaky.androidide.plugins.aiassistant.memory
2+
3+
import android.app.ActivityManager
4+
import android.content.Context
5+
6+
/**
7+
* Free RAM, as an interface so the pre-flight check is testable without a device.
8+
*/
9+
fun interface DeviceMemory {
10+
11+
/**
12+
* Must be read at the moment of the check, never cached — the user may have just closed apps.
13+
*
14+
* @return free RAM in bytes, or null when it cannot be read. Null rather than a negative
15+
* sentinel, so "unknown" cannot be confused with the reading of 0 a device under real
16+
* pressure will genuinely report.
17+
*/
18+
fun availableBytes(): Long?
19+
}
20+
21+
/**
22+
* Reads free RAM from [ActivityManager], which needs no permission.
23+
*
24+
* @param contextProvider supplies the Android context, and null before the plugin is initialized
25+
* @param onReadError reports why a reading could not be taken. The pre-flight only learns that one
26+
* could not be, so without this the reason behind a silently skipped check is lost.
27+
*/
28+
class SystemDeviceMemory(
29+
private val contextProvider: () -> Context?,
30+
private val onReadError: (Throwable) -> Unit = {},
31+
) : DeviceMemory {
32+
33+
/** Returns null rather than throwing; the gate treats "unknown" as "don't warn". */
34+
override fun availableBytes(): Long? = try {
35+
// applicationContext, so a provider that ever hands back an Activity can't be held here.
36+
val context = contextProvider()?.applicationContext ?: return null
37+
val manager = context.getSystemService(Context.ACTIVITY_SERVICE) as ActivityManager
38+
ActivityManager.MemoryInfo().also(manager::getMemoryInfo).availMem
39+
} catch (e: Exception) {
40+
onReadError(e)
41+
null
42+
}
43+
}

0 commit comments

Comments
 (0)