diff --git a/ai-assistant/src/main/assets/docs/index.html b/ai-assistant/src/main/assets/docs/index.html
index 3b69e634..e0799942 100644
--- a/ai-assistant/src/main/assets/docs/index.html
+++ b/ai-assistant/src/main/assets/docs/index.html
@@ -56,6 +56,28 @@ Choosing a backend
agent reads are sent to Google over HTTPS.
+ Will this model fit in memory?
+ When you pick a local .gguf file, it is measured against the RAM
+ free on the device at that moment, before anything is loaded. If it looks too
+ large you get a warning with the actual figures and two choices:
+
+ - Cancel — nothing is saved and nothing is loaded. Whatever model you
+ had selected before stays in use. Dismissing the warning with Back does the
+ same thing.
+ - Proceed anyway — the model is accepted despite the shortfall. Use
+ this when you are about to close other apps, for example.
+
+ The warning quotes two numbers. Memory to load is the weights: these are
+ memory-mapped, so they need not all fit at once — when they don't, the device pages
+ them in and out, which is why an oversized model can stall for minutes instead of
+ failing immediately. Memory to run is the KV cache and compute buffers, which
+ are ordinary allocations and do have to fit. That is why a model can be reported as
+ risky rather than impossible: the outcome genuinely depends on how much paging your
+ device will tolerate.
+ To fit a large model, close other apps and select it again, or pick a smaller or
+ more heavily quantized build. A Q4_K_M quantization of a 1–3B model is the most
+ likely to run comfortably.
+
What the agent can do
- Read and search files within the current project.
diff --git a/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/AiAssistantPlugin.kt b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/AiAssistantPlugin.kt
index 0b5b73d2..604eb87e 100644
--- a/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/AiAssistantPlugin.kt
+++ b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/AiAssistantPlugin.kt
@@ -66,6 +66,10 @@ class AiAssistantPlugin : IPlugin, UIExtension, DocumentationExtension, Settings
const val TOOLTIP_TAG_SETTINGS_GEMINI_MODEL = "ai_settings_gemini_model"
const val TOOLTIP_TAG_SETTINGS_GET_KEY = "ai_settings_get_free_key"
+ // Tags for the memory pre-flight warning (see MemoryWarningDialogFragment).
+ const val TOOLTIP_TAG_MEMORY_PROCEED = "agent_memory_warning_proceed"
+ const val TOOLTIP_TAG_MEMORY_CANCEL = "agent_memory_warning_cancel"
+
@Volatile
private var pluginContext: PluginContext? = null
@@ -457,6 +461,39 @@ class AiAssistantPlugin : IPlugin, UIExtension, DocumentationExtension, Settings
models can't generate replies. Larger models are slower and use
more memory; the file is copied into the app's private storage on
first use.
+ The model is measured against this device's free memory before it
+ is accepted. If it looks too large you get the figures and a choice
+ to cancel or continue.
+ """.trimIndent(),
+ buttons = listOf(
+ PluginTooltipButton(description = "AI Assistant guide", uri = "index.html", order = 0)
+ )
+ ),
+ PluginTooltipEntry(
+ tag = TOOLTIP_TAG_MEMORY_PROCEED,
+ summary = "Load this model anyway, accepting that it may fail or slow the device.",
+ detail = """
+ The model's weights plus its working memory look larger than the
+ RAM free right now. Weights are memory-mapped, so a load can still
+ succeed by paging — which is why the outcome is a risk rather than a
+ certainty: it may work, fail quickly, or stall for minutes first.
+ Use this when you know the numbers are wrong for your situation,
+ for example because you are about to close other apps.
+ """.trimIndent(),
+ buttons = listOf(
+ PluginTooltipButton(description = "AI Assistant guide", uri = "index.html", order = 0)
+ )
+ ),
+ PluginTooltipEntry(
+ tag = TOOLTIP_TAG_MEMORY_CANCEL,
+ summary = "Abandon this model; the previously selected one is left untouched.",
+ detail = """
+ Nothing is saved and nothing is loaded, so the model you had
+ selected before stays in use. This is the safe choice, and also what
+ happens if you dismiss the warning with Back.
+ To fit a large model, close other apps and pick it again, or
+ choose a smaller or more heavily quantized build — a Q4_K_M
+ quantization of a 1–3B model is the most likely to run.
""".trimIndent(),
buttons = listOf(
PluginTooltipButton(description = "AI Assistant guide", uri = "index.html", order = 0)
diff --git a/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/fragments/AiSettingsFragment.kt b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/fragments/AiSettingsFragment.kt
index 3f734839..a6b5c064 100644
--- a/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/fragments/AiSettingsFragment.kt
+++ b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/fragments/AiSettingsFragment.kt
@@ -17,8 +17,10 @@ import android.widget.*
import androidx.activity.result.contract.ActivityResultContracts
import androidx.annotation.DrawableRes
import androidx.fragment.app.Fragment
+import androidx.lifecycle.Lifecycle
import androidx.lifecycle.ViewModelProvider
import androidx.lifecycle.lifecycleScope
+import androidx.lifecycle.repeatOnLifecycle
import com.google.android.material.dialog.MaterialAlertDialogBuilder
import com.itsaky.androidide.plugins.PluginContext
import com.itsaky.androidide.plugins.aiassistant.AiAssistantPlugin
@@ -31,6 +33,7 @@ import com.itsaky.androidide.plugins.aiassistant.viewmodel.AiBackend
import com.itsaky.androidide.plugins.aiassistant.viewmodel.AiSettingsViewModel
import com.itsaky.androidide.plugins.aiassistant.viewmodel.EngineState
import com.itsaky.androidide.plugins.aiassistant.viewmodel.ModelLoadingState
+import com.itsaky.androidide.plugins.aiassistant.viewmodel.ModelMemoryWarning
import kotlinx.coroutines.launch
import java.text.SimpleDateFormat
import java.util.Date
@@ -42,7 +45,7 @@ import kotlin.math.roundToInt
* chat's own shortcuts. The host mounts it full-screen in PluginScreenActivity, which provides no
* toolbar, so this fragment brings its own app bar and closes by finishing that activity.
*/
-class AiSettingsFragment : Fragment() {
+class AiSettingsFragment : Fragment(), MemoryWarningDialogFragment.Host {
private lateinit var viewModel: AiSettingsViewModel
private lateinit var settingsToolbar: LinearLayout
@@ -130,6 +133,57 @@ class AiSettingsFragment : Fragment() {
initializeViews(view)
setupToolbar()
setupBackendSelector()
+ observeMemoryWarnings()
+ }
+
+ /**
+ * Puts a "this model may not fit" question to the user. Collected under STARTED so the dialog is
+ * never shown to a stopped fragment; the event waits in the ViewModel until then.
+ */
+ private fun observeMemoryWarnings() {
+ dropStaleMemoryWarning()
+ viewLifecycleOwner.lifecycleScope.launch {
+ viewLifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
+ viewModel.modelMemoryWarnings.collect(::showMemoryWarning)
+ }
+ }
+ }
+
+ /**
+ * Dismiss a warning dialog the framework restored around a question that no longer exists.
+ * After process death the load that raised it is gone, so every button on it would be a silent
+ * no-op — better to take it away than to leave the user pressing a dialog that decides nothing.
+ */
+ private fun dropStaleMemoryWarning() {
+ if (viewModel.hasPendingMemoryWarning) return
+ val restored = childFragmentManager.findFragmentByTag(MemoryWarningDialogFragment.TAG)
+ (restored as? MemoryWarningDialogFragment)?.dismissAllowingStateLoss()
+ }
+
+ /**
+ * Shown as a child fragment, so it survives rotation and can still reach this host. Must stay
+ * idempotent: an unanswered question is re-published to every new collector by
+ * [com.itsaky.androidide.plugins.aiassistant.viewmodel.UserConfirmation].
+ *
+ * @param warning the model and the figures to put to the user
+ */
+ private fun showMemoryWarning(warning: ModelMemoryWarning) {
+ if (childFragmentManager.findFragmentByTag(MemoryWarningDialogFragment.TAG) != null) return
+ MemoryWarningDialogFragment.newInstance(warning)
+ .show(childFragmentManager, MemoryWarningDialogFragment.TAG)
+ }
+
+ override fun onModelMemoryDecision(proceed: Boolean) {
+ viewModel.onMemoryWarningDecision(proceed)
+ // Not requireContext(): onCancel can reach us as the fragment is going away.
+ val ctx = context ?: return
+ if (!proceed) {
+ Toast.makeText(
+ ctx,
+ getString(R.string.llm_memory_warning_declined),
+ Toast.LENGTH_LONG,
+ ).show()
+ }
}
override fun onResume() {
diff --git a/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/fragments/MemoryWarningDialogFragment.kt b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/fragments/MemoryWarningDialogFragment.kt
new file mode 100644
index 00000000..1a204d10
--- /dev/null
+++ b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/fragments/MemoryWarningDialogFragment.kt
@@ -0,0 +1,153 @@
+package com.itsaky.androidide.plugins.aiassistant.fragments
+
+import android.app.Dialog
+import android.content.DialogInterface
+import android.os.Bundle
+import android.view.View
+import androidx.fragment.app.DialogFragment
+import com.google.android.material.dialog.MaterialAlertDialogBuilder
+import com.itsaky.androidide.plugins.aiassistant.AiAssistantPlugin
+import com.itsaky.androidide.plugins.aiassistant.R
+import com.itsaky.androidide.plugins.aiassistant.memory.ModelMemoryGate
+import com.itsaky.androidide.plugins.aiassistant.util.ByteSize
+import com.itsaky.androidide.plugins.aiassistant.viewmodel.ModelMemoryWarning
+import com.itsaky.androidide.plugins.base.PluginFragmentHelper
+import com.itsaky.androidide.plugins.services.IdeTooltipService
+
+/**
+ * Warns that a selected model may be too large for this device, offering **Cancel** — the default
+ * in every sense visible to the user — or **Proceed anyway** (ADFA-1798). Decisions go to [Host],
+ * resolved per call so they still arrive after the framework recreates this on a rotation.
+ */
+class MemoryWarningDialogFragment : DialogFragment() {
+
+ private val tooltipService: IdeTooltipService? by lazy {
+ try {
+ PluginFragmentHelper.getServiceRegistry(AiAssistantPlugin.PLUGIN_ID)
+ ?.get(IdeTooltipService::class.java)
+ } catch (e: Exception) {
+ // Tooltip help is optional; long-press simply shows nothing when it's unavailable.
+ AiAssistantPlugin.getContext()?.logger
+ ?.warn("MemoryWarningDialogFragment: tooltip service unavailable", e)
+ null
+ }
+ }
+
+ /**
+ * Receives this dialog's outcome. Implemented by the fragment that shows the dialog, which must
+ * be its **parent** fragment (show it with `childFragmentManager`).
+ */
+ interface Host {
+ /** @param proceed true to load the model anyway, false to abandon the selection. */
+ fun onModelMemoryDecision(proceed: Boolean)
+ }
+
+ companion object {
+ const val TAG = "ModelMemoryWarning"
+
+ private const val ARG_MODEL_NAME = "model_name"
+ private const val ARG_LOAD_BYTES = "load_bytes"
+ private const val ARG_RUN_BYTES = "run_bytes"
+ private const val ARG_AVAILABLE_BYTES = "available_bytes"
+ private const val ARG_SEVERITY = "severity"
+
+ /**
+ * Builds the dialog. Everything it needs is in [getArguments], so the framework can recreate
+ * it after a configuration change while the answer channel stays on the ViewModel.
+ *
+ * @param warning the model and the figures to show
+ */
+ fun newInstance(warning: ModelMemoryWarning): MemoryWarningDialogFragment =
+ MemoryWarningDialogFragment().apply {
+ arguments = Bundle().apply {
+ putString(ARG_MODEL_NAME, warning.modelName)
+ putLong(ARG_LOAD_BYTES, warning.loadBytes)
+ putLong(ARG_RUN_BYTES, warning.runBytes)
+ putLong(ARG_AVAILABLE_BYTES, warning.availableBytes)
+ putString(ARG_SEVERITY, warning.severity.name)
+ }
+ }
+ }
+
+ override fun onCreateDialog(savedInstanceState: Bundle?): Dialog {
+ val args = arguments ?: Bundle.EMPTY
+ val messageId = when (severity()) {
+ ModelMemoryGate.Severity.INSUFFICIENT -> R.string.llm_memory_warning_insufficient
+ ModelMemoryGate.Severity.TIGHT -> R.string.llm_memory_warning_tight
+ }
+ val message = getString(
+ messageId,
+ args.getString(ARG_MODEL_NAME).orEmpty(),
+ ByteSize.format(args.getLong(ARG_LOAD_BYTES)),
+ ByteSize.format(args.getLong(ARG_RUN_BYTES)),
+ ByteSize.format(args.getLong(ARG_AVAILABLE_BYTES)),
+ )
+
+ val dialog = MaterialAlertDialogBuilder(requireContext())
+ .setTitle(getString(R.string.llm_memory_warning_title))
+ .setMessage(message)
+ // Cancel takes the positive slot: Material emphasizes that button.
+ .setPositiveButton(getString(R.string.llm_memory_warning_cancel)) { _, _ -> decide(false) }
+ .setNegativeButton(getString(R.string.llm_memory_warning_proceed)) { _, _ -> decide(true) }
+ .create()
+
+ // Bound after show(): the buttons don't exist before it.
+ dialog.setOnShowListener {
+ val cancelButton = dialog.getButton(Dialog.BUTTON_POSITIVE)
+ wireTooltip(cancelButton, AiAssistantPlugin.TOOLTIP_TAG_MEMORY_CANCEL)
+ wireTooltip(
+ dialog.getButton(Dialog.BUTTON_NEGATIVE),
+ AiAssistantPlugin.TOOLTIP_TAG_MEMORY_PROCEED,
+ )
+ // So Enter, D-pad and a screen reader all reach the safe choice first.
+ cancelButton?.requestFocus()
+ }
+
+ return dialog
+ }
+
+ /**
+ * Back press and outside tap both mean "don't risk it". Overridden rather than set on the
+ * builder: [DialogFragment] replaces a builder cancel listener in `prepareDialog`, so the
+ * awaiting caller would never learn the model was declined.
+ *
+ * @param dialog the dialog that was cancelled
+ */
+ override fun onCancel(dialog: DialogInterface) {
+ super.onCancel(dialog)
+ decide(false)
+ }
+
+ /** Falls back to the milder wording, so a bad argument can't overstate the risk. */
+ private fun severity(): ModelMemoryGate.Severity = try {
+ ModelMemoryGate.Severity.valueOf(
+ arguments?.getString(ARG_SEVERITY) ?: ModelMemoryGate.Severity.TIGHT.name
+ )
+ } catch (e: IllegalArgumentException) {
+ ModelMemoryGate.Severity.TIGHT
+ }
+
+ /**
+ * Delivers the outcome to the host fragment. Resolved on each call rather than captured at
+ * construction, so it still works on the instance the framework recreated after a rotation.
+ */
+ private fun decide(proceed: Boolean) {
+ val host = parentFragment as? Host
+ if (host == null) {
+ // Nothing else reports this, and the load that raised the warning waits forever on it.
+ AiAssistantPlugin.getContext()?.logger
+ ?.warn("MemoryWarningDialogFragment: no Host parent; decision dropped")
+ return
+ }
+ host.onModelMemoryDecision(proceed)
+ }
+
+ /** Shows this plugin's tooltip for [tag] when [view] is long-pressed (Tier 1/2 + guide). */
+ private fun wireTooltip(view: View?, tag: String) {
+ view?.setOnLongClickListener { anchor ->
+ val service = tooltipService ?: return@setOnLongClickListener false
+ service.showTooltip(anchor, AiAssistantPlugin.TOOLTIP_CATEGORY, tag)
+ true
+ }
+ }
+}
diff --git a/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/memory/DeviceMemory.kt b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/memory/DeviceMemory.kt
new file mode 100644
index 00000000..52a1ebf8
--- /dev/null
+++ b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/memory/DeviceMemory.kt
@@ -0,0 +1,43 @@
+package com.itsaky.androidide.plugins.aiassistant.memory
+
+import android.app.ActivityManager
+import android.content.Context
+
+/**
+ * Free RAM, as an interface so the pre-flight check is testable without a device.
+ */
+fun interface DeviceMemory {
+
+ /**
+ * Must be read at the moment of the check, never cached — the user may have just closed apps.
+ *
+ * @return free RAM in bytes, or null when it cannot be read. Null rather than a negative
+ * sentinel, so "unknown" cannot be confused with the reading of 0 a device under real
+ * pressure will genuinely report.
+ */
+ fun availableBytes(): Long?
+}
+
+/**
+ * Reads free RAM from [ActivityManager], which needs no permission.
+ *
+ * @param contextProvider supplies the Android context, and null before the plugin is initialized
+ * @param onReadError reports why a reading could not be taken. The pre-flight only learns that one
+ * could not be, so without this the reason behind a silently skipped check is lost.
+ */
+class SystemDeviceMemory(
+ private val contextProvider: () -> Context?,
+ private val onReadError: (Throwable) -> Unit = {},
+) : DeviceMemory {
+
+ /** Returns null rather than throwing; the gate treats "unknown" as "don't warn". */
+ override fun availableBytes(): Long? = try {
+ // applicationContext, so a provider that ever hands back an Activity can't be held here.
+ val context = contextProvider()?.applicationContext ?: return null
+ val manager = context.getSystemService(Context.ACTIVITY_SERVICE) as ActivityManager
+ ActivityManager.MemoryInfo().also(manager::getMemoryInfo).availMem
+ } catch (e: Exception) {
+ onReadError(e)
+ null
+ }
+}
diff --git a/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/memory/ModelMemoryEstimator.kt b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/memory/ModelMemoryEstimator.kt
new file mode 100644
index 00000000..9e010b4e
--- /dev/null
+++ b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/memory/ModelMemoryEstimator.kt
@@ -0,0 +1,108 @@
+package com.itsaky.androidide.plugins.aiassistant.memory
+
+import com.itsaky.androidide.plugins.aiassistant.util.GgufHeader
+
+/**
+ * What a model will cost in memory, split the way it behaves at runtime.
+ *
+ * @property loadBytes the weights. mmap'd, so they need not *fit*: when they don't, the device
+ * thrashes page cache instead of failing fast, which is the "ten minutes, then an error" report.
+ * @property runBytes KV cache and compute buffers. Ordinary allocations, so this part must fit.
+ * @property fromHeader true when [runBytes] came from the model's own shape values rather than the
+ * size-based fallback; diagnostics only.
+ */
+data class MemoryEstimate(
+ val loadBytes: Long,
+ val runBytes: Long,
+ val fromHeader: Boolean,
+) {
+ /** Saturating: [loadBytes] is a queried file size, and a bogus one must not wrap into "it fits". */
+ val totalBytes: Long
+ get() = if (loadBytes > Long.MAX_VALUE - runBytes) Long.MAX_VALUE else loadBytes + runBytes
+}
+
+/**
+ * Estimates the memory a `.gguf` model needs, from its size and its declared shape. Pure and
+ * Android-free, so the arithmetic is unit-testable. The context and batch sizes below are ai-core's,
+ * fixed on its native side: an estimate has to model the loader that will actually run.
+ */
+object ModelMemoryEstimator {
+
+ /**
+ * The context every load gets, hard-coded as `ctx_params.n_ctx` in ai-core's `llama-android.cpp`.
+ * The KV cache is sized from it, so keep the two in step.
+ */
+ const val RUNTIME_CONTEXT_TOKENS = 4096L
+
+ /** Two bytes per cached element: f16, the default KV type. */
+ private const val KV_BYTES_PER_ELEMENT = 2L
+
+ /**
+ * Graph and compute buffers for the 2048-token batch ai-core allocates (`new_batch` in
+ * `LLamaAndroid.load`). Not derivable from the header, so it is a flat allowance.
+ */
+ private const val COMPUTE_BUFFER_BYTES = 256L * 1024 * 1024
+
+ /**
+ * Floor for the size-based fallback, matching `ModelLoadDiagnostics.MIN_RUN_BYTES` in ai-core so
+ * this warning and the refusal that gates the load itself cannot contradict each other.
+ */
+ private const val MIN_FALLBACK_RUN_BYTES = 256L * 1024 * 1024
+
+ /**
+ * Ceilings on the header's shape values, each far above the largest real model. They exist so a
+ * corrupt or crafted file cannot wrap the KV-cache product: within them it stays below 2^57, and
+ * an out-of-range value falls back to the size-based heuristic instead of a wrong estimate.
+ */
+ private const val MAX_LAYERS = 1L shl 10
+ private const val MAX_HEADS = 1L shl 12
+ private const val MAX_WIDTH = 1L shl 20
+
+ /**
+ * @param fileSizeBytes the model file's size, or null when it is unknown
+ * @param header the model's metadata, or null when it could not be read
+ * @return the estimate, or null when there is nothing to base one on
+ */
+ fun estimate(fileSizeBytes: Long?, header: GgufHeader?): MemoryEstimate? {
+ if (fileSizeBytes == null || fileSizeBytes <= 0L) return null
+ val kvCacheBytes = header?.let(::kvCacheBytes)
+ return if (kvCacheBytes != null) {
+ MemoryEstimate(fileSizeBytes, kvCacheBytes + COMPUTE_BUFFER_BYTES, fromHeader = true)
+ } else {
+ // A quarter of the weights is a rough stand-in for a cache we couldn't measure.
+ MemoryEstimate(
+ loadBytes = fileSizeBytes,
+ runBytes = maxOf(MIN_FALLBACK_RUN_BYTES, fileSizeBytes / 4),
+ fromHeader = false,
+ )
+ }
+ }
+
+ /**
+ * KV cache size for a full context: one key and one value entry per kv head, per layer, per
+ * position. Null unless every value it needs is present and within its ceiling.
+ */
+ private fun kvCacheBytes(header: GgufHeader): Long? {
+ val layers = header.blockCount?.within(MAX_LAYERS) ?: return null
+ val heads = header.headCount?.within(MAX_HEADS) ?: return null
+ // Grouped-query attention caches only the kv heads; absent means one per head (plain MHA).
+ val kvHeads = (header.headCountKv ?: heads).within(MAX_HEADS) ?: return null
+ val keyWidth = header.keyLength?.within(MAX_WIDTH) ?: defaultHeadWidth(header) ?: return null
+ val valueWidth = header.valueLength?.within(MAX_WIDTH) ?: defaultHeadWidth(header) ?: return null
+ return KV_BYTES_PER_ELEMENT * layers * RUNTIME_CONTEXT_TOKENS * kvHeads * (keyWidth + valueWidth)
+ }
+
+ /** The value when it is positive and no larger than [ceiling]; null when it is neither. */
+ private fun Long.within(ceiling: Long): Long? = takeIf { it > 0L && it <= ceiling }
+
+ /**
+ * Head width where the model does not state one: the width split evenly across the heads. Only
+ * a default — gemma-3, for one, declares a key/value width that is not this quotient — so it is
+ * used solely for the models that leave `attention.key_length` and `.value_length` out.
+ */
+ private fun defaultHeadWidth(header: GgufHeader): Long? {
+ val embedding = header.embeddingLength?.within(MAX_WIDTH) ?: return null
+ val heads = header.headCount?.within(MAX_HEADS) ?: return null
+ return (embedding / heads).within(MAX_WIDTH)
+ }
+}
diff --git a/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/memory/ModelMemoryGate.kt b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/memory/ModelMemoryGate.kt
new file mode 100644
index 00000000..a62ad500
--- /dev/null
+++ b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/memory/ModelMemoryGate.kt
@@ -0,0 +1,53 @@
+package com.itsaky.androidide.plugins.aiassistant.memory
+
+/**
+ * Decides whether a selected model is worth warning about, from its estimate and the free RAM.
+ *
+ * Pure, so every boundary is unit-testable. It never refuses a model: the caller asks the user,
+ * who may proceed anyway (ADFA-1798).
+ */
+object ModelMemoryGate {
+
+ /** How badly the device is short, which is the difference between "may" and "will" fail. */
+ enum class Severity {
+ /** The runtime allocations fit, but not alongside the weights: expect thrashing. */
+ TIGHT,
+
+ /** Not even the KV cache and compute buffers fit: expect a quick, hard failure. */
+ INSUFFICIENT,
+ }
+
+ sealed interface Verdict {
+
+ /** Everything fits; load without interrupting the user. */
+ data object Safe : Verdict
+
+ /** Nothing to judge — the size or the free RAM was unreadable. Fails open: no warning. */
+ data object Unknown : Verdict
+
+ /**
+ * @param estimate what the model is expected to need
+ * @param availableBytes free RAM at the moment of the check
+ * @param severity how short the device is
+ */
+ data class Risky(
+ val estimate: MemoryEstimate,
+ val availableBytes: Long,
+ val severity: Severity,
+ ) : Verdict
+ }
+
+ /**
+ * @param estimate the model's expected cost, or null when it could not be estimated
+ * @param availableBytes free RAM, or null when it could not be read
+ * @return the verdict for this model on this device, right now
+ */
+ fun evaluate(estimate: MemoryEstimate?, availableBytes: Long?): Verdict {
+ if (estimate == null || availableBytes == null) return Verdict.Unknown
+ return when {
+ availableBytes >= estimate.totalBytes -> Verdict.Safe
+ availableBytes >= estimate.runBytes -> Verdict.Risky(estimate, availableBytes, Severity.TIGHT)
+ else -> Verdict.Risky(estimate, availableBytes, Severity.INSUFFICIENT)
+ }
+ }
+}
diff --git a/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/util/ByteSize.kt b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/util/ByteSize.kt
new file mode 100644
index 00000000..cf89ba5a
--- /dev/null
+++ b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/util/ByteSize.kt
@@ -0,0 +1,25 @@
+package com.itsaky.androidide.plugins.aiassistant.util
+
+import java.util.Locale
+
+/**
+ * Formats byte counts for display, in binary units and the US locale. Mirrors ai-core's `ByteSize`,
+ * unreachable from here across isolated plugin classloaders — keep the units in step, since both
+ * describe the same memory and mixing binary with decimal would read as a contradiction.
+ */
+internal object ByteSize {
+
+ private const val BYTES_PER_MB = 1024.0 * 1024.0
+ private const val BYTES_PER_GB = BYTES_PER_MB * 1024.0
+
+ /**
+ * Formats [bytes] with the largest unit that keeps the figure meaningful; sub-gigabyte values
+ * stay in MB, since "0.3 GB free" reads as a broken string rather than as a shortage.
+ *
+ * @param bytes a byte count
+ * @return the size as a one-decimal "X.X GB" or "X.X MB" string
+ */
+ fun format(bytes: Long): String =
+ if (bytes >= BYTES_PER_GB) String.format(Locale.US, "%.1f GB", bytes / BYTES_PER_GB)
+ else String.format(Locale.US, "%.1f MB", bytes / BYTES_PER_MB)
+}
diff --git a/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/util/GgufHeaderReader.kt b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/util/GgufHeaderReader.kt
new file mode 100644
index 00000000..02288064
--- /dev/null
+++ b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/util/GgufHeaderReader.kt
@@ -0,0 +1,305 @@
+package com.itsaky.androidide.plugins.aiassistant.util
+
+import java.io.BufferedInputStream
+import java.io.DataInputStream
+import java.io.EOFException
+import java.io.FilterInputStream
+import java.io.IOException
+import java.io.InputStream
+
+/**
+ * The GGUF metadata the memory estimate needs. Every field is nullable because a file may omit any
+ * key; the estimator then falls back to a heuristic instead of guessing.
+ *
+ * @property architecture `general.architecture`; also the prefix every other key here is read under
+ * @property blockCount transformer layers, `{arch}.block_count`
+ * @property embeddingLength model width, `{arch}.embedding_length`
+ * @property headCount attention heads, `{arch}.attention.head_count`
+ * @property headCountKv key/value heads under grouped-query attention, absent for plain MHA
+ * @property keyLength per-head key width, `{arch}.attention.key_length`; absent means it is the
+ * model width divided by the head count, which is only the default and not always the truth
+ * @property valueLength per-head value width, `{arch}.attention.value_length`; as [keyLength]
+ */
+data class GgufHeader(
+ val architecture: String?,
+ val blockCount: Long?,
+ val embeddingLength: Long?,
+ val headCount: Long?,
+ val headCountKv: Long?,
+ val keyLength: Long? = null,
+ val valueLength: Long? = null,
+)
+
+/**
+ * Reads the metadata block at the front of a `.gguf` file — the shape values the KV-cache estimate
+ * needs. Never reads the weights, and fails closed to null: an unreadable header must mean "no
+ * estimate", never a wrong one. Mirrors ai-core's `GgufModelInspector` — keep the two in sync.
+ */
+internal object GgufHeaderReader {
+
+ /** "GGUF", little-endian. */
+ private const val GGUF_MAGIC = 0x46554747
+
+ // GGUF metadata value types.
+ private const val T_UINT8 = 0
+ private const val T_INT8 = 1
+ private const val T_UINT16 = 2
+ private const val T_INT16 = 3
+ private const val T_UINT32 = 4
+ private const val T_INT32 = 5
+ private const val T_FLOAT32 = 6
+ private const val T_BOOL = 7
+ private const val T_STRING = 8
+ private const val T_ARRAY = 9
+ private const val T_UINT64 = 10
+ private const val T_INT64 = 11
+ private const val T_FLOAT64 = 12
+
+ private const val KEY_ARCHITECTURE = "general.architecture"
+
+ // Matched by suffix, then attributed to the "{arch}." prefix they carry — see [readHeader].
+ private const val SUFFIX_BLOCK_COUNT = ".block_count"
+ private const val SUFFIX_EMBEDDING_LENGTH = ".embedding_length"
+ private const val SUFFIX_HEAD_COUNT = ".attention.head_count"
+ private const val SUFFIX_HEAD_COUNT_KV = ".attention.head_count_kv"
+ private const val SUFFIX_KEY_LENGTH = ".attention.key_length"
+ private const val SUFFIX_VALUE_LENGTH = ".attention.value_length"
+
+ /** Backstop against a corrupt count claiming millions of entries; real files hold dozens. */
+ private const val MAX_METADATA_ENTRIES = 4096L
+
+ /**
+ * Backstop on one array's declared length. The byte budget alone bounds the data read but not
+ * the iteration count: an array of 1-byte elements would spin ~67M times before tripping it.
+ * Far above the largest real value here, a ~256k-token tokenizer vocabulary.
+ */
+ private const val MAX_ARRAY_ELEMENTS = 1L shl 22
+
+ /** Keys and architecture names are tiny; anything longer means a corrupt length field. */
+ private const val MAX_STRING_BYTES = 1L shl 20
+
+ /**
+ * Ceiling on the bytes one parse may consume. Real metadata blocks run to a few MB, so this is
+ * generous — it exists because an array's element count is otherwise unbounded, and a corrupt
+ * one would have the parse skip its way through a multi-gigabyte model file before giving up.
+ */
+ private const val MAX_METADATA_BYTES = 64L shl 20
+
+ /**
+ * Blocking I/O, and reads at most [MAX_METADATA_BYTES] — never call this on the main thread.
+ *
+ * @param openStream opens the candidate model, or returns null when it can't be opened
+ * @return the metadata values that were present, or null if the header could not be parsed
+ */
+ fun read(openStream: () -> InputStream?): GgufHeader? = try {
+ openStream()?.use { stream ->
+ val budgeted = BudgetedInputStream(stream, MAX_METADATA_BYTES)
+ readHeader(DataInputStream(BufferedInputStream(budgeted, 1 shl 16)))
+ }
+ } catch (e: Throwable) {
+ // Throwable, so a StackOverflowError is a null header rather than a crash.
+ null
+ }
+
+ /**
+ * Aborts the parse once [limit] bytes have been consumed, so no declared count can make the
+ * read run on past the metadata. Throwing is deliberate: [read] treats it like any other parse
+ * failure and returns null, which puts the estimate on its size-based fallback.
+ */
+ private class BudgetedInputStream(source: InputStream, private val limit: Long) :
+ FilterInputStream(source) {
+
+ private var consumed = 0L
+
+ override fun read(): Int = super.read().also { if (it >= 0) charge(1L) }
+
+ override fun read(b: ByteArray, off: Int, len: Int): Int =
+ super.read(b, off, len).also { if (it > 0) charge(it.toLong()) }
+
+ /**
+ * Clamped to the remaining budget *before* delegating, since [charge] only runs once the
+ * delegate returns: one string declaring a length of 2^62 would otherwise read its way
+ * through the whole model file, 2 KB at a time, with the budget never consulted.
+ */
+ override fun skip(n: Long): Long {
+ if (n <= 0L) return 0L
+ // +1 so the clamp itself can still push consumed past the limit and trip charge().
+ val allowed = minOf(n, limit - consumed + 1)
+ if (allowed <= 0L) throw IOException("GGUF metadata exceeded $limit bytes")
+ return super.skip(allowed).also { if (it > 0) charge(it) }
+ }
+
+ private fun charge(bytes: Long) {
+ consumed += bytes
+ if (consumed > limit) throw IOException("GGUF metadata exceeded $limit bytes")
+ }
+ }
+
+ /** The shape values seen under one `{arch}.` prefix. A file may carry more than one. */
+ private class ArchShape {
+ var blockCount: Long? = null
+ var embeddingLength: Long? = null
+ var headCount: Long? = null
+ var headCountKv: Long? = null
+ var keyLength: Long? = null
+ var valueLength: Long? = null
+ }
+
+ private fun readHeader(input: DataInputStream): GgufHeader? {
+ if (readU32(input) != GGUF_MAGIC) return null
+
+ // v1 used 32-bit counts and string lengths; v2+ use 64-bit.
+ val wide = readU32(input) >= 2
+ readCount(input, wide) // tensor count, unused here
+ val entryCount = readCount(input, wide)
+ // Give up rather than truncate: a head_count_kv never reached reads as plain MHA.
+ if (entryCount < 0L || entryCount > MAX_METADATA_ENTRIES) return null
+
+ var architecture: String? = null
+ // Keyed by each key's own "{arch}", so a multimodal file's `clip.…` values stay out.
+ val shapes = HashMap()
+
+ // Walked to the end: an absent key is meaningful, and only the end of the block proves it.
+ var index = 0L
+ while (index < entryCount) {
+ val key = readString(input, wide)
+ val type = readU32(input)
+ when {
+ key == KEY_ARCHITECTURE && type == T_STRING -> architecture = readString(input, wide)
+
+ key.endsWith(SUFFIX_BLOCK_COUNT) ->
+ shapeFor(shapes, key, SUFFIX_BLOCK_COUNT).blockCount = readInteger(input, type, wide)
+
+ key.endsWith(SUFFIX_EMBEDDING_LENGTH) ->
+ shapeFor(shapes, key, SUFFIX_EMBEDDING_LENGTH).embeddingLength = readInteger(input, type, wide)
+
+ // Before the head_count suffix, so the more specific key always wins.
+ key.endsWith(SUFFIX_HEAD_COUNT_KV) ->
+ shapeFor(shapes, key, SUFFIX_HEAD_COUNT_KV).headCountKv = readInteger(input, type, wide)
+
+ key.endsWith(SUFFIX_HEAD_COUNT) ->
+ shapeFor(shapes, key, SUFFIX_HEAD_COUNT).headCount = readInteger(input, type, wide)
+
+ key.endsWith(SUFFIX_KEY_LENGTH) ->
+ shapeFor(shapes, key, SUFFIX_KEY_LENGTH).keyLength = readInteger(input, type, wide)
+
+ key.endsWith(SUFFIX_VALUE_LENGTH) ->
+ shapeFor(shapes, key, SUFFIX_VALUE_LENGTH).valueLength = readInteger(input, type, wide)
+
+ else -> skipValue(input, type, wide)
+ }
+ index++
+ }
+
+ val shape = architecture?.let(shapes::get)
+ return GgufHeader(
+ architecture = architecture,
+ blockCount = shape?.blockCount,
+ embeddingLength = shape?.embeddingLength,
+ headCount = shape?.headCount,
+ headCountKv = shape?.headCountKv,
+ keyLength = shape?.keyLength,
+ valueLength = shape?.valueLength,
+ )
+ }
+
+ /** The shape values for the architecture [key] belongs to, created on first sight. */
+ private fun shapeFor(shapes: HashMap, key: String, suffix: String): ArchShape =
+ shapes.getOrPut(key.dropLast(suffix.length)) { ArchShape() }
+
+ /**
+ * Consumes one value, returning it when it is an integer and null otherwise. The value is
+ * always consumed either way, or the next key would be read from the middle of it.
+ */
+ private fun readInteger(input: DataInputStream, type: Int, wide: Boolean): Long? = when (type) {
+ T_UINT8, T_INT8 -> input.readUnsignedByte().toLong()
+ T_UINT16, T_INT16 -> readU16(input)
+ T_UINT32, T_INT32 -> readU32(input).toLong() and 0xFFFFFFFFL
+ T_UINT64, T_INT64 -> readU64(input)
+ else -> {
+ skipValue(input, type, wide)
+ null
+ }
+ }
+
+ private fun skipValue(input: DataInputStream, type: Int, wide: Boolean) {
+ when (type) {
+ T_UINT8, T_INT8, T_BOOL -> skipFully(input, 1)
+ T_UINT16, T_INT16 -> skipFully(input, 2)
+ T_UINT32, T_INT32, T_FLOAT32 -> skipFully(input, 4)
+ T_UINT64, T_INT64, T_FLOAT64 -> skipFully(input, 8)
+ T_STRING -> skipFully(input, readCount(input, wide))
+ T_ARRAY -> {
+ val elementType = readU32(input)
+ // The format forbids nesting, and following one recurses until the stack goes.
+ if (elementType == T_ARRAY) throw IllegalStateException("Nested GGUF array")
+ val elements = readCount(input, wide)
+ if (elements < 0L || elements > MAX_ARRAY_ELEMENTS) {
+ throw IllegalStateException("Unreasonable GGUF array length: $elements")
+ }
+ var i = 0L
+ while (i < elements) {
+ skipValue(input, elementType, wide)
+ i++
+ }
+ }
+ else -> throw IllegalStateException("Unknown GGUF value type: $type")
+ }
+ }
+
+ // --- little-endian primitives ---
+
+ private fun readU16(input: DataInputStream): Long {
+ val b0 = input.read()
+ val b1 = input.read()
+ if (b1 < 0) throw EOFException()
+ return ((b0 and 0xFF) or ((b1 and 0xFF) shl 8)).toLong()
+ }
+
+ private fun readU32(input: DataInputStream): Int {
+ val b0 = input.read()
+ val b1 = input.read()
+ val b2 = input.read()
+ val b3 = input.read()
+ if (b3 < 0) throw EOFException()
+ return (b0 and 0xFF) or ((b1 and 0xFF) shl 8) or ((b2 and 0xFF) shl 16) or ((b3 and 0xFF) shl 24)
+ }
+
+ private fun readU64(input: DataInputStream): Long {
+ var value = 0L
+ for (i in 0 until 8) {
+ val b = input.read()
+ if (b < 0) throw EOFException()
+ value = value or ((b.toLong() and 0xFF) shl (8 * i))
+ }
+ return value
+ }
+
+ /** A length or count field: 64-bit on GGUF v2+, 32-bit on v1. */
+ private fun readCount(input: DataInputStream, wide: Boolean): Long =
+ if (wide) readU64(input) else readU32(input).toLong() and 0xFFFFFFFFL
+
+ private fun readString(input: DataInputStream, wide: Boolean): String {
+ val length = readCount(input, wide)
+ if (length < 0 || length > MAX_STRING_BYTES) {
+ throw IllegalStateException("Unreasonable GGUF string length: $length")
+ }
+ val bytes = ByteArray(length.toInt())
+ input.readFully(bytes)
+ return String(bytes, Charsets.UTF_8)
+ }
+
+ private fun skipFully(input: DataInputStream, count: Long) {
+ var remaining = count
+ while (remaining > 0) {
+ val skipped = input.skip(remaining)
+ if (skipped > 0) {
+ remaining -= skipped
+ } else {
+ // skip() can return 0 near buffer boundaries; fall back to a read.
+ if (input.read() < 0) throw EOFException()
+ remaining--
+ }
+ }
+ }
+}
diff --git a/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/util/ModelFileSource.kt b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/util/ModelFileSource.kt
new file mode 100644
index 00000000..86c07ba8
--- /dev/null
+++ b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/util/ModelFileSource.kt
@@ -0,0 +1,132 @@
+package com.itsaky.androidide.plugins.aiassistant.util
+
+import android.content.Context
+import android.content.Intent
+import android.net.Uri
+import android.provider.OpenableColumns
+import java.io.File
+import java.io.InputStream
+
+/**
+ * What a selected model file is called and how big it is.
+ *
+ * @property displayName the name to show the user; never blank
+ * @property sizeBytes the file's size, or null when it could not be established — the pre-flight
+ * then has nothing to weigh and skips the check rather than guessing
+ */
+data class ModelFileInfo(val displayName: String, val sizeBytes: Long?)
+
+/**
+ * Reads a selected model file's metadata and bytes, whether it came from the document picker as a
+ * `content://` URI or from a saved filesystem path. An interface, so the memory pre-flight can be
+ * exercised against ordinary files without a device.
+ */
+interface ModelFileSource {
+
+ /**
+ * Name and size together — for a content URI that is one provider query, where asking
+ * separately costs two IPC round trips for one row. Do NOT call on the main thread.
+ *
+ * @param context supplies the resolver that holds the picker's permission grant
+ * @param uriString the selected model, as a `content://` URI or a filesystem path
+ */
+ fun info(context: Context, uriString: String): ModelFileInfo
+
+ /** Opens the model for reading; null when it cannot be opened. Not for the main thread. */
+ fun openStream(context: Context, uriString: String): InputStream?
+
+ /** Decoded last path segment — a cheap name that at least avoids raw `%3A` escapes. */
+ fun fallbackDisplayName(uriOrPath: String): String
+
+ /**
+ * Give back the persistable read grant the picker took for [uriString], for a model the user
+ * ended up not keeping — the grant table has a hard per-app limit. A no-op for a filesystem
+ * path, and for a grant that was never held.
+ */
+ fun releaseAccess(context: Context, uriString: String)
+}
+
+/**
+ * [ModelFileSource] over the document provider and the filesystem.
+ *
+ * Every lookup degrades rather than throwing: an unnamed file falls back to its path, and an
+ * unknown size is reported as unknown. A model the user picked is not a place to fail hard.
+ *
+ * @param onError reports a failed lookup, so a silently skipped pre-flight can still be explained
+ */
+class ContentModelFileSource(
+ private val onError: (String, Throwable) -> Unit = { _, _ -> },
+) : ModelFileSource {
+
+ override fun info(context: Context, uriString: String): ModelFileInfo =
+ if (uriString.startsWith(CONTENT_SCHEME)) {
+ documentInfo(context, uriString) ?: ModelFileInfo(fallbackDisplayName(uriString), null)
+ } else {
+ ModelFileInfo(fallbackDisplayName(uriString), fileSize(uriString))
+ }
+
+ override fun openStream(context: Context, uriString: String): InputStream? = try {
+ if (uriString.startsWith(CONTENT_SCHEME)) {
+ context.contentResolver.openInputStream(Uri.parse(uriString))
+ } else {
+ File(uriString).takeIf { it.isFile }?.inputStream()
+ }
+ } catch (e: Exception) {
+ onError("could not open $uriString", e)
+ null
+ }
+
+ override fun fallbackDisplayName(uriOrPath: String): String =
+ (try {
+ Uri.decode(uriOrPath)
+ } catch (e: Exception) {
+ uriOrPath
+ }).substringAfterLast('/')
+
+ override fun releaseAccess(context: Context, uriString: String) {
+ if (!uriString.startsWith(CONTENT_SCHEME)) return
+ try {
+ context.contentResolver.releasePersistableUriPermission(
+ Uri.parse(uriString),
+ Intent.FLAG_GRANT_READ_URI_PERMISSION,
+ )
+ } catch (e: Exception) {
+ // Never held, or already released — nothing is broken either way.
+ onError("could not release the read grant for $uriString", e)
+ }
+ }
+
+ /** One query for both columns; null when the provider answered with neither. */
+ private fun documentInfo(context: Context, uriString: String): ModelFileInfo? = try {
+ context.contentResolver
+ .query(Uri.parse(uriString), arrayOf(NAME_COLUMN, SIZE_COLUMN), null, null, null)
+ ?.use { cursor ->
+ if (!cursor.moveToFirst()) return@use null
+ val nameIndex = cursor.getColumnIndex(NAME_COLUMN)
+ val sizeIndex = cursor.getColumnIndex(SIZE_COLUMN)
+ val name = nameIndex.takeIf { it >= 0 && !cursor.isNull(it) }
+ ?.let { cursor.getString(it) }
+ ?.takeIf { it.isNotBlank() }
+ val size = sizeIndex.takeIf { it >= 0 && !cursor.isNull(it) }
+ ?.let { cursor.getLong(it) }
+ ?.takeIf { it > 0L }
+ ModelFileInfo(name ?: fallbackDisplayName(uriString), size)
+ }
+ } catch (e: Exception) {
+ onError("could not read the metadata of $uriString", e)
+ null
+ }
+
+ private fun fileSize(path: String): Long? = try {
+ File(path).length().takeIf { it > 0L }
+ } catch (e: Exception) {
+ onError("could not read the size of $path", e)
+ null
+ }
+
+ private companion object {
+ const val CONTENT_SCHEME = "content://"
+ val NAME_COLUMN: String = OpenableColumns.DISPLAY_NAME
+ val SIZE_COLUMN: String = OpenableColumns.SIZE
+ }
+}
diff --git a/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/viewmodel/AiSettingsViewModel.kt b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/viewmodel/AiSettingsViewModel.kt
index b6a7ae46..0f5149b2 100644
--- a/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/viewmodel/AiSettingsViewModel.kt
+++ b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/viewmodel/AiSettingsViewModel.kt
@@ -10,12 +10,22 @@ import com.itsaky.androidide.plugins.aiassistant.gemini.GeminiCatalogGateway
import com.itsaky.androidide.plugins.aiassistant.gemini.KeyVerification
import com.itsaky.androidide.plugins.aiassistant.gemini.ReflectiveGeminiCatalogGateway
import com.itsaky.androidide.plugins.aiassistant.gemini.toKeyVerification
+import com.itsaky.androidide.plugins.aiassistant.memory.DeviceMemory
+import com.itsaky.androidide.plugins.aiassistant.memory.ModelMemoryEstimator
+import com.itsaky.androidide.plugins.aiassistant.memory.ModelMemoryGate
+import com.itsaky.androidide.plugins.aiassistant.memory.SystemDeviceMemory
import com.itsaky.androidide.plugins.aiassistant.security.SecureApiKeyStore
import com.itsaky.androidide.plugins.aiassistant.R
+import com.itsaky.androidide.plugins.aiassistant.util.ByteSize
+import com.itsaky.androidide.plugins.aiassistant.util.ContentModelFileSource
import com.itsaky.androidide.plugins.aiassistant.util.GgufFileInspector
+import com.itsaky.androidide.plugins.aiassistant.util.GgufHeaderReader
+import com.itsaky.androidide.plugins.aiassistant.util.ModelFileInfo
+import com.itsaky.androidide.plugins.aiassistant.util.ModelFileSource
import kotlinx.coroutines.CancellationException
import kotlinx.coroutines.CoroutineDispatcher
import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.launch
import kotlinx.coroutines.withContext
import com.itsaky.androidide.plugins.PluginContext
@@ -58,12 +68,45 @@ enum class AiBackend(val displayName: String) {
*/
data class GeminiModelOptions(val models: List, val isLive: Boolean)
+/**
+ * A selected model that may not fit in this device's memory, with the figures to show the user.
+ *
+ * @param modelName the model's display name
+ * @param loadBytes memory the weights need
+ * @param runBytes memory the KV cache and compute buffers need on top of the weights
+ * @param availableBytes free RAM when the check ran
+ * @param severity whether the shortfall makes failure likely or merely possible
+ */
+data class ModelMemoryWarning(
+ val modelName: String,
+ val loadBytes: Long,
+ val runBytes: Long,
+ val availableBytes: Long,
+ val severity: ModelMemoryGate.Severity,
+)
+
+/**
+ * @param deviceMemory free-RAM reading for the pre-flight; null builds the live one, which cannot
+ * be a default argument because it needs [logger], and a default cannot reach an instance member
+ * @param modelFiles reads a selected model's name, size and bytes; null builds the live one
+ */
class AiSettingsViewModel(
private val getContext: () -> PluginContext?,
private val ioDispatcher: CoroutineDispatcher = Dispatchers.IO,
- private val catalogGateway: GeminiCatalogGateway = ReflectiveGeminiCatalogGateway()
+ private val catalogGateway: GeminiCatalogGateway = ReflectiveGeminiCatalogGateway(),
+ deviceMemory: DeviceMemory? = null,
+ modelFiles: ModelFileSource? = null,
) : ViewModel() {
+ private val deviceMemory: DeviceMemory = deviceMemory ?: SystemDeviceMemory(
+ contextProvider = { getContext()?.androidContext },
+ onReadError = { e -> logger?.warn("$TAG: could not read free memory", e) },
+ )
+
+ private val modelFiles: ModelFileSource = modelFiles ?: ContentModelFileSource { what, e ->
+ logger?.warn("$TAG: $what", e)
+ }
+
companion object {
private const val TAG = "AiSettingsViewModel"
@@ -94,6 +137,22 @@ class AiSettingsViewModel(
private val _engineState = MutableLiveData(EngineState.Initialized)
val engineState: LiveData get() = _engineState
+ /** The memory pre-flight's consent gate; see [loadModelFromUri]. */
+ private val memoryConfirmation = UserConfirmation()
+
+ /**
+ * Models that may not fit in memory, to be put to the user as a warning. One-shot events: each
+ * is delivered once, and the answer comes back through [onMemoryWarningDecision].
+ */
+ val modelMemoryWarnings: Flow get() = memoryConfirmation.requests
+
+ /**
+ * Whether a memory warning is actually waiting on an answer. False after process death, where
+ * the dialog is restored by the framework but the load that raised it is long gone — the UI
+ * uses this to drop a dialog whose answer nobody would receive.
+ */
+ val hasPendingMemoryWarning: Boolean get() = memoryConfirmation.hasOutstandingRequest
+
init {
checkInitialState()
}
@@ -105,14 +164,21 @@ class AiSettingsViewModel(
// For plugin, engine is always "ready" since it's managed by ai-core plugin
_engineState.value = EngineState.Initialized
- // Reflect a previously selected model so it survives closing/reopening settings,
- // using the display name persisted at load time (no content-provider query here).
- _modelLoadingState.value = if (savedPath != null) {
+ _modelLoadingState.value = modelStateFor(savedPath)
+ }
+
+ /**
+ * The state describing the model that is actually configured. Built from the name persisted at
+ * load time, so it needs no provider query.
+ *
+ * @param savedPath the stored model path, or null when none is configured
+ */
+ private fun modelStateFor(savedPath: String?): ModelLoadingState =
+ if (savedPath != null) {
ModelLoadingState.Loaded(getSavedModelName() ?: fallbackDisplayName(savedPath))
} else {
ModelLoadingState.Idle
}
- }
/**
* This plugin's settings store — and, for the Gemini keys, ai-core's too.
@@ -138,33 +204,7 @@ class AiSettingsViewModel(
}
/** Decoded last path segment — a cheap fallback that at least avoids raw %3A escapes. */
- fun fallbackDisplayName(uriOrPath: String): String =
- (try { android.net.Uri.decode(uriOrPath) } catch (e: Exception) { uriOrPath }).substringAfterLast('/')
-
- /**
- * Resolve the real file name for a selected model. For a `content://` URI this queries the
- * document provider's [OpenableColumns.DISPLAY_NAME] (e.g. "Llama-3.2-1B.gguf"); otherwise,
- * and on any failure, it falls back to the decoded last path segment. Do NOT call on the main
- * thread — the provider query can block.
- */
- private fun resolveDisplayName(uriString: String): String {
- if (uriString.startsWith("content://")) {
- try {
- val uri = android.net.Uri.parse(uriString)
- getContext()?.androidContext?.contentResolver
- ?.query(uri, arrayOf(android.provider.OpenableColumns.DISPLAY_NAME), null, null, null)
- ?.use { c ->
- val idx = c.getColumnIndex(android.provider.OpenableColumns.DISPLAY_NAME)
- if (idx >= 0 && c.moveToFirst() && !c.isNull(idx)) {
- c.getString(idx)?.takeIf { it.isNotBlank() }?.let { return it }
- }
- }
- } catch (e: Exception) {
- logger?.warn("$TAG: could not resolve display name for $uriString", e)
- }
- }
- return fallbackDisplayName(uriString)
- }
+ fun fallbackDisplayName(uriOrPath: String): String = modelFiles.fallbackDisplayName(uriOrPath)
fun getAvailableBackends(): List = AiBackend.entries
@@ -264,7 +304,7 @@ class AiSettingsViewModel(
*/
suspend fun saveGeminiApiKey(apiKey: String, verified: Boolean = false): Boolean =
withContext(ioDispatcher) {
- // Checked first: returning true here would have the UI claim an unwritten key was saved.
+ // Checked first, or the UI would claim an unwritten key was saved.
val prefs = getPluginPrefs()
if (prefs == null) {
logger?.error("$TAG: cannot save Gemini API key: plugin preferences unavailable")
@@ -389,19 +429,23 @@ class AiSettingsViewModel(
}
/**
- * Load a model from URI.
- * In the plugin context, we just save the path - the actual loading
- * is handled by the ai-core plugin's LocalLlmBackend.
+ * Saves the selected model's path; ai-core's `LocalLlmBackend` does the loading itself. That
+ * write is what makes it load, so the memory pre-flight gates it: a model the user declines is
+ * never stored, and therefore never loaded (ADFA-1798).
+ *
+ * @param uriString the selected model, as a `content://` URI or a filesystem path
+ * @param context resolves the model's display name, size and header
*/
fun loadModelFromUri(uriString: String, context: Context) {
- viewModelScope.launch(Dispatchers.IO) {
+ viewModelScope.launch(ioDispatcher) {
_modelLoadingState.postValue(ModelLoadingState.Loading)
try {
- // Resolve the real file name (not the raw content-URI doc id) for display.
- val fileName = resolveDisplayName(uriString)
+ // One lookup for both: the real file name to show, and the size to estimate from.
+ val fileInfo = modelFiles.info(context, uriString)
+ val fileName = fileInfo.displayName
- // Reject a non-GGUF pick up front, so no bad path is persisted or shown as "Loaded".
+ // Rejected up front, so no bad path is persisted or shown as "Loaded".
if (!GgufFileInspector.looksLikeGguf(context.contentResolver, uriString)) {
_modelLoadingState.postValue(
ModelLoadingState.Error(context.getString(R.string.error_model_not_gguf, fileName))
@@ -409,17 +453,28 @@ class AiSettingsViewModel(
return@launch
}
+ if (!confirmMemoryHeadroom(uriString, fileInfo, context)) {
+ logger?.info("$TAG: model declined at the memory warning: $fileName")
+ // Never the configured model: re-checking it and declining must not revoke it.
+ if (uriString != getLocalModelPath()) {
+ modelFiles.releaseAccess(context, uriString)
+ }
+ restoreSavedModelState()
+ return@launch
+ }
+
// Persist the name before the path so the savedModelPath observer can read it.
saveLocalModelName(fileName)
saveLocalModelPath(uriString)
- // In plugin context, we don't directly load the model
- // The ai-core plugin will load it when needed
+ // Nothing is loaded here; ai-core reads this path when it needs the model.
_modelLoadingState.postValue(
ModelLoadingState.Loaded(fileName)
)
logger?.debug("$TAG: model path saved: $uriString ($fileName)")
+ } catch (e: CancellationException) {
+ throw e
} catch (e: Exception) {
logger?.error("$TAG: error saving model path", e)
_modelLoadingState.postValue(
@@ -428,4 +483,70 @@ class AiSettingsViewModel(
}
}
}
+
+ /**
+ * Answers an outstanding [modelMemoryWarnings] question. Safe to call from the main thread, and
+ * a no-op when nothing is waiting.
+ *
+ * @param proceed true to load the model anyway, false to abandon the selection
+ */
+ fun onMemoryWarningDecision(proceed: Boolean) {
+ memoryConfirmation.answer(proceed)
+ }
+
+ /**
+ * Checks the model against free RAM and, when it looks too large, asks the user whether to go
+ * ahead. Fails OPEN: an unreadable size or header means no warning rather than a wrong one.
+ *
+ * @return true to continue with this model
+ */
+ private suspend fun confirmMemoryHeadroom(
+ uriString: String,
+ fileInfo: ModelFileInfo,
+ context: Context
+ ): Boolean {
+ val modelName = fileInfo.displayName
+ val estimate = ModelMemoryEstimator.estimate(
+ fileSizeBytes = fileInfo.sizeBytes,
+ header = GgufHeaderReader.read { modelFiles.openStream(context, uriString) },
+ )
+ // Read last and never cached: the user may have just closed apps to make room.
+ val availableBytes = deviceMemory.availableBytes()
+
+ return when (val verdict = ModelMemoryGate.evaluate(estimate, availableBytes)) {
+ ModelMemoryGate.Verdict.Safe -> true
+
+ ModelMemoryGate.Verdict.Unknown -> {
+ val missing = if (estimate == null) "the model's size" else "free memory"
+ logger?.warn("$TAG: could not read $missing; skipping the pre-flight for $modelName")
+ true
+ }
+
+ is ModelMemoryGate.Verdict.Risky -> {
+ logger?.warn(
+ "$TAG: $modelName may not fit: needs ${ByteSize.format(verdict.estimate.loadBytes)}" +
+ " + ${ByteSize.format(verdict.estimate.runBytes)} to run," +
+ " ${ByteSize.format(verdict.availableBytes)} free (${verdict.severity})"
+ )
+ memoryConfirmation.ask(
+ ModelMemoryWarning(
+ modelName = modelName,
+ loadBytes = verdict.estimate.loadBytes,
+ runBytes = verdict.estimate.runBytes,
+ availableBytes = verdict.availableBytes,
+ severity = verdict.severity,
+ )
+ )
+ }
+ }
+ }
+
+ /**
+ * Republishes the model that is actually configured, so abandoning a selection leaves the
+ * screen describing the previous model rather than the one that was never stored.
+ */
+ private fun restoreSavedModelState() {
+ _modelLoadingState.postValue(modelStateFor(getLocalModelPath()))
+ }
+
}
diff --git a/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/viewmodel/UserConfirmation.kt b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/viewmodel/UserConfirmation.kt
new file mode 100644
index 00000000..01556000
--- /dev/null
+++ b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/viewmodel/UserConfirmation.kt
@@ -0,0 +1,68 @@
+package com.itsaky.androidide.plugins.aiassistant.viewmodel
+
+import kotlinx.coroutines.CompletableDeferred
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.filterNotNull
+import kotlinx.coroutines.sync.Mutex
+import kotlinx.coroutines.sync.withLock
+
+/**
+ * A yes/no question for the user that a coroutine can await, so a flow needing consent stays one
+ * readable sequence instead of being cut in half by a callback. Holds no Android types. Thread-safe,
+ * and one question at a time: a second [ask] waits for the first to be answered.
+ */
+internal class UserConfirmation {
+
+ // State, not an event: a one-shot stream drops the question if its collector dies mid-delivery.
+ private val outstanding = MutableStateFlow(null)
+
+ /**
+ * The question to put to the user, re-emitted to each new collector until it is answered — so
+ * the UI must be idempotent (see `AiSettingsFragment.showMemoryWarning`).
+ */
+ val requests: Flow = outstanding.filterNotNull()
+
+ /**
+ * Whether a question is currently waiting on the user. False after process death, where the UI
+ * may have been restored around a question this object no longer knows anything about — see
+ * `AiSettingsFragment.dropStaleMemoryWarning`.
+ */
+ val hasOutstandingRequest: Boolean get() = outstanding.value != null
+
+ /** Serializes questions, so [pending] can only ever describe one of them. */
+ private val askMutex = Mutex()
+
+ @Volatile
+ private var pending: CompletableDeferred? = null
+
+ /**
+ * Publishes [request] and suspends until the UI answers. Cancelling the caller (the ViewModel
+ * being cleared, say) abandons the question rather than resolving it either way.
+ *
+ * @param request what to ask
+ * @return the user's answer
+ */
+ suspend fun ask(request: T): Boolean = askMutex.withLock {
+ val answer = CompletableDeferred()
+ pending = answer
+ try {
+ outstanding.value = request
+ answer.await()
+ } finally {
+ pending = null
+ // Withdraws the question, so the next collector is not asked one nobody is waiting on.
+ outstanding.value = null
+ }
+ }
+
+ /**
+ * Answers the outstanding question; a no-op when there is none, or when it already has an
+ * answer, so a duplicate reply from a recreated dialog is harmless.
+ *
+ * @param granted true to proceed, false to decline
+ */
+ fun answer(granted: Boolean) {
+ pending?.complete(granted)
+ }
+}
diff --git a/ai-assistant/src/main/res/values/strings.xml b/ai-assistant/src/main/res/values/strings.xml
index da079f21..0ca39851 100644
--- a/ai-assistant/src/main/res/values/strings.xml
+++ b/ai-assistant/src/main/res/values/strings.xml
@@ -199,4 +199,12 @@
\"%1$s\" isn\'t a valid .gguf model (it may be the wrong file or a corrupt or partial download). Select a .gguf chat model.
+
+
+ This model may be too large
+ \"%1$s\" needs about %2$s of memory to load, plus about %3$s more to run. This device has about %4$s available right now, which may not be enough. It might work, or fail quickly, or fail after several minutes.\n\nClose other apps to free memory, or choose a smaller or more heavily quantized model (for example a Q4_K_M build of a 1–3B model).
+ \"%1$s\" needs about %2$s of memory to load, plus about %3$s more to run. This device has about %4$s available right now, which is not enough. Loading it will most likely fail, and may make the IDE unresponsive first.\n\nClose other apps to free memory, or choose a smaller or more heavily quantized model (for example a Q4_K_M build of a 1–3B model).
+ Proceed anyway
+ Cancel
+ Model not selected. The previously selected model, if any, is unchanged.
diff --git a/ai-assistant/src/test/kotlin/com/itsaky/androidide/plugins/aiassistant/memory/ModelMemoryEstimatorTest.kt b/ai-assistant/src/test/kotlin/com/itsaky/androidide/plugins/aiassistant/memory/ModelMemoryEstimatorTest.kt
new file mode 100644
index 00000000..a1c4ac61
--- /dev/null
+++ b/ai-assistant/src/test/kotlin/com/itsaky/androidide/plugins/aiassistant/memory/ModelMemoryEstimatorTest.kt
@@ -0,0 +1,180 @@
+package com.itsaky.androidide.plugins.aiassistant.memory
+
+import com.itsaky.androidide.plugins.aiassistant.util.GgufHeader
+import org.junit.Assert.assertEquals
+import org.junit.Assert.assertFalse
+import org.junit.Assert.assertNull
+import org.junit.Assert.assertTrue
+import org.junit.Test
+
+/**
+ * Tests for the memory arithmetic behind the pre-flight warning.
+ *
+ * The figures reach the user, so they are asserted exactly rather than as ranges: a wrong KV-cache
+ * formula would still produce a plausible-looking dialog.
+ */
+class ModelMemoryEstimatorTest {
+
+ private val megabyte = 1024L * 1024
+ private val computeBuffer = 256 * megabyte
+
+ /** gemma-3-1b: 26 layers, 1152 wide, 4 heads, 1 kv head — so a 288-wide kv projection. */
+ private val gemma = GgufHeader(
+ architecture = "gemma3",
+ blockCount = 26,
+ embeddingLength = 1152,
+ headCount = 4,
+ headCountKv = 1,
+ )
+
+ @Test
+ fun givenAModelShape_whenEstimated_thenTheKvCacheIsSizedForTheFullContext() {
+ val fileSize = 800 * megabyte
+ // A literal: restating the formula asserts only that the code agrees with itself.
+ val expectedKvCache = 122_683_392L
+
+ val estimate = ModelMemoryEstimator.estimate(fileSize, gemma)
+
+ assertEquals(fileSize, estimate?.loadBytes)
+ assertEquals(expectedKvCache + computeBuffer, estimate?.runBytes)
+ assertTrue(estimate?.fromHeader == true)
+ }
+
+ @Test
+ fun givenAModelWithoutGroupedQueryAttention_whenEstimated_thenEveryHeadIsCached() {
+ val mha = gemma.copy(headCountKv = null)
+
+ val estimate = ModelMemoryEstimator.estimate(800 * megabyte, mha)
+
+ // One kv head per attention head: four times gemma's 122,683,392-byte cache.
+ assertEquals(490_733_568L + computeBuffer, estimate?.runBytes)
+ }
+
+ @Test
+ fun givenADeclaredKeyAndValueWidth_whenEstimated_thenTheyAreUsedOverTheHeadQuotient() {
+ // gemma-3 declares 256, not the 288 embedding / heads implies — a 12% overstatement.
+ val declared = gemma.copy(keyLength = 256, valueLength = 256)
+
+ val estimate = ModelMemoryEstimator.estimate(800 * megabyte, declared)
+
+ // 2 x 26 x 4096 x 1 x (256 + 256).
+ assertEquals(109_051_904L + computeBuffer, estimate?.runBytes)
+ }
+
+ @Test
+ fun givenOnlyOneOfTheTwoWidthsDeclared_whenEstimated_thenTheOtherStillFallsBack() {
+ val halfDeclared = gemma.copy(keyLength = 256)
+
+ val estimate = ModelMemoryEstimator.estimate(800 * megabyte, halfDeclared)
+
+ // 2 x 26 x 4096 x 1 x (256 declared key + 288 derived value).
+ assertEquals(115_867_648L + computeBuffer, estimate?.runBytes)
+ }
+
+ @Test
+ fun givenAZeroedKeyWidth_whenEstimated_thenTheDerivedWidthIsUsedInstead() {
+ val corrupt = gemma.copy(keyLength = 0, valueLength = 0)
+
+ val estimate = ModelMemoryEstimator.estimate(800 * megabyte, corrupt)
+
+ assertEquals(122_683_392L + computeBuffer, estimate?.runBytes)
+ }
+
+ @Test
+ fun givenNoHeader_whenEstimated_thenItFallsBackToAShareOfTheFileSize() {
+ val fileSize = 4096L * megabyte
+
+ val estimate = ModelMemoryEstimator.estimate(fileSize, header = null)
+
+ assertEquals(fileSize, estimate?.loadBytes)
+ assertEquals(fileSize / 4, estimate?.runBytes)
+ assertFalse(estimate?.fromHeader == true)
+ }
+
+ @Test
+ fun givenASmallModelAndNoHeader_whenEstimated_thenTheRuntimeFloorStillApplies() {
+ // A quarter of a 300 MB file is nowhere near enough for a KV cache at full context.
+ val estimate = ModelMemoryEstimator.estimate(300 * megabyte, header = null)
+
+ assertEquals(256 * megabyte, estimate?.runBytes)
+ }
+
+ @Test
+ fun givenAnIncompleteHeader_whenEstimated_thenItFallsBackRatherThanGuessing() {
+ val partial = gemma.copy(embeddingLength = null)
+
+ val estimate = ModelMemoryEstimator.estimate(4096L * megabyte, partial)
+
+ assertEquals(1024 * megabyte, estimate?.runBytes)
+ assertFalse(estimate?.fromHeader == true)
+ }
+
+ @Test
+ fun givenAZeroedHeaderValue_whenEstimated_thenItFallsBackInsteadOfDividingByZero() {
+ val corrupt = gemma.copy(headCount = 0)
+
+ val estimate = ModelMemoryEstimator.estimate(800 * megabyte, corrupt)
+
+ assertEquals(256 * megabyte, estimate?.runBytes)
+ assertFalse(estimate?.fromHeader == true)
+ }
+
+ @Test
+ fun givenALayerCountThatWouldWrapTheProduct_whenEstimated_thenItFallsBackInsteadOfUnderstating() {
+ // 2^61 layers with 288-wide projections wraps the KV term to exactly zero.
+ val crafted = gemma.copy(blockCount = 1L shl 61, keyLength = 288, valueLength = 288)
+
+ val estimate = ModelMemoryEstimator.estimate(800 * megabyte, crafted)
+
+ assertEquals(256 * megabyte, estimate?.runBytes)
+ assertFalse(estimate?.fromHeader == true)
+ }
+
+ @Test
+ fun givenALayerCountThatWouldMakeTheProductNegative_whenEstimated_thenTheTotalStaysPositive() {
+ // 2^49 layers with 1-wide projections lands on Long.MIN_VALUE, which reads as "it fits".
+ val crafted = gemma.copy(blockCount = 1L shl 49, keyLength = 1, valueLength = 1)
+
+ val estimate = ModelMemoryEstimator.estimate(800 * megabyte, crafted)!!
+
+ assertTrue(estimate.totalBytes > 0L)
+ assertFalse(estimate.fromHeader)
+ }
+
+ @Test
+ fun givenTheLargestRealisticShape_whenEstimated_thenTheCeilingsStillAcceptIt() {
+ // llama-3.1-405B: 126 layers, 16384 wide, 128 heads, 8 kv heads.
+ val large = GgufHeader(
+ architecture = "llama",
+ blockCount = 126,
+ embeddingLength = 16384,
+ headCount = 128,
+ headCountKv = 8,
+ )
+
+ val estimate = ModelMemoryEstimator.estimate(200L * 1024 * megabyte, large)
+
+ assertTrue(estimate?.fromHeader == true)
+ }
+
+ @Test
+ fun givenAFileSizeThatWouldWrapTheTotal_whenTotalled_thenItSaturatesInsteadOfGoingNegative() {
+ // A queried SIZE column is not trustworthy, and a negative total reads as "it fits".
+ val estimate = ModelMemoryEstimator.estimate(Long.MAX_VALUE, gemma)!!
+
+ assertEquals(Long.MAX_VALUE, estimate.totalBytes)
+ }
+
+ @Test
+ fun givenAnUnknownFileSize_whenEstimated_thenThereIsNoEstimate() {
+ assertNull(ModelMemoryEstimator.estimate(null, gemma))
+ assertNull(ModelMemoryEstimator.estimate(0L, gemma))
+ }
+
+ @Test
+ fun givenAnEstimate_whenTotalled_thenItIsTheWeightsPlusTheRuntime() {
+ val estimate = ModelMemoryEstimator.estimate(800 * megabyte, gemma)!!
+
+ assertEquals(estimate.loadBytes + estimate.runBytes, estimate.totalBytes)
+ }
+}
diff --git a/ai-assistant/src/test/kotlin/com/itsaky/androidide/plugins/aiassistant/memory/ModelMemoryGateTest.kt b/ai-assistant/src/test/kotlin/com/itsaky/androidide/plugins/aiassistant/memory/ModelMemoryGateTest.kt
new file mode 100644
index 00000000..02c1d60b
--- /dev/null
+++ b/ai-assistant/src/test/kotlin/com/itsaky/androidide/plugins/aiassistant/memory/ModelMemoryGateTest.kt
@@ -0,0 +1,79 @@
+package com.itsaky.androidide.plugins.aiassistant.memory
+
+import com.itsaky.androidide.plugins.aiassistant.memory.ModelMemoryGate.Severity
+import com.itsaky.androidide.plugins.aiassistant.memory.ModelMemoryGate.Verdict
+import org.junit.Assert.assertEquals
+import org.junit.Test
+
+/**
+ * Tests for the decision to warn, including its boundaries — one byte either side of them is the
+ * difference between an interrupted user and a crash.
+ */
+class ModelMemoryGateTest {
+
+ private val megabyte = 1024L * 1024
+
+ /** 800 MB of weights that need 400 MB of working memory: 1200 MB in total. */
+ private val estimate = MemoryEstimate(
+ loadBytes = 800 * megabyte,
+ runBytes = 400 * megabyte,
+ fromHeader = true,
+ )
+
+ private fun severityAt(availableBytes: Long): Severity? =
+ (ModelMemoryGate.evaluate(estimate, availableBytes) as? Verdict.Risky)?.severity
+
+ @Test
+ fun givenRoomForEverything_whenEvaluated_thenTheUserIsNotInterrupted() {
+ assertEquals(Verdict.Safe, ModelMemoryGate.evaluate(estimate, 2048 * megabyte))
+ }
+
+ @Test
+ fun givenExactlyEnough_whenEvaluated_thenItIsSafe() {
+ assertEquals(Verdict.Safe, ModelMemoryGate.evaluate(estimate, estimate.totalBytes))
+ }
+
+ @Test
+ fun givenOneByteTooLittle_whenEvaluated_thenThrashingIsTheRisk() {
+ assertEquals(Severity.TIGHT, severityAt(estimate.totalBytes - 1))
+ }
+
+ @Test
+ fun givenRoomForTheRuntimeButNotTheWeights_whenEvaluated_thenThrashingIsTheRisk() {
+ assertEquals(Severity.TIGHT, severityAt(600 * megabyte))
+ }
+
+ @Test
+ fun givenNotEvenRoomForTheRuntime_whenEvaluated_thenFailureIsExpected() {
+ assertEquals(Severity.INSUFFICIENT, severityAt(estimate.runBytes - 1))
+ }
+
+ @Test
+ fun givenExactlyEnoughForTheRuntime_whenEvaluated_thenItIsOnlyTight() {
+ assertEquals(Severity.TIGHT, severityAt(estimate.runBytes))
+ }
+
+ @Test
+ fun givenNoEstimate_whenEvaluated_thenNothingIsClaimed() {
+ assertEquals(Verdict.Unknown, ModelMemoryGate.evaluate(null, 2048 * megabyte))
+ }
+
+ @Test
+ fun givenUnreadableMemory_whenEvaluated_thenItFailsOpenRatherThanWarning() {
+ // A device we can't measure must not be told its model won't fit.
+ assertEquals(Verdict.Unknown, ModelMemoryGate.evaluate(estimate, null))
+ }
+
+ @Test
+ fun givenNoFreeMemoryAtAll_whenEvaluated_thenThatIsARealReadingAndNotUnknown() {
+ assertEquals(Severity.INSUFFICIENT, severityAt(0L))
+ }
+
+ @Test
+ fun givenARiskyModel_whenEvaluated_thenTheVerdictCarriesTheFiguresToShow() {
+ val verdict = ModelMemoryGate.evaluate(estimate, 600 * megabyte) as Verdict.Risky
+
+ assertEquals(estimate, verdict.estimate)
+ assertEquals(600 * megabyte, verdict.availableBytes)
+ }
+}
diff --git a/ai-assistant/src/test/kotlin/com/itsaky/androidide/plugins/aiassistant/util/GgufHeaderReaderTest.kt b/ai-assistant/src/test/kotlin/com/itsaky/androidide/plugins/aiassistant/util/GgufHeaderReaderTest.kt
new file mode 100644
index 00000000..26088ed7
--- /dev/null
+++ b/ai-assistant/src/test/kotlin/com/itsaky/androidide/plugins/aiassistant/util/GgufHeaderReaderTest.kt
@@ -0,0 +1,283 @@
+package com.itsaky.androidide.plugins.aiassistant.util
+
+import org.junit.Assert.assertEquals
+import org.junit.Assert.assertNull
+import org.junit.Test
+import java.io.ByteArrayInputStream
+import java.io.IOException
+import java.io.InputStream
+
+/**
+ * Tests for the GGUF metadata read that the memory estimate is based on.
+ *
+ * Every case runs against real bytes from [GgufWriter], since the point of the reader is that it
+ * agrees with the on-disk format — and that it gives up cleanly when it doesn't.
+ */
+class GgufHeaderReaderTest {
+
+ private fun read(bytes: ByteArray): GgufHeader? =
+ GgufHeaderReader.read { ByteArrayInputStream(bytes) }
+
+ /** The shape of gemma-3-1b, as a representative grouped-query-attention model. */
+ private fun gemmaHeader(version: Int = 3) = GgufWriter(version)
+ .string("general.architecture", "gemma3")
+ .string("general.name", "gemma-3-1b-it")
+ .uint32("gemma3.block_count", 26)
+ .uint32("gemma3.embedding_length", 1152)
+ .uint32("gemma3.attention.head_count", 4)
+ .uint32("gemma3.attention.head_count_kv", 1)
+ .build()
+
+ @Test
+ fun givenAGgufHeader_whenRead_thenEveryShapeValueIsReturned() {
+ val header = read(gemmaHeader())
+
+ assertEquals("gemma3", header?.architecture)
+ assertEquals(26L, header?.blockCount)
+ assertEquals(1152L, header?.embeddingLength)
+ assertEquals(4L, header?.headCount)
+ assertEquals(1L, header?.headCountKv)
+ }
+
+ @Test
+ fun givenAVersionOneHeader_whenRead_thenItsNarrowerLengthsAreUnderstood() {
+ // v1 wrote 32-bit counts and string lengths; misreading them desynchronizes every field.
+ val header = read(gemmaHeader(version = 1))
+
+ assertEquals("gemma3", header?.architecture)
+ assertEquals(26L, header?.blockCount)
+ assertEquals(1L, header?.headCountKv)
+ }
+
+ @Test
+ fun givenAModelWithoutGroupedQueryAttention_whenRead_thenTheKvHeadCountIsAbsent() {
+ val bytes = GgufWriter()
+ .string("general.architecture", "llama")
+ .uint32("llama.block_count", 32)
+ .uint32("llama.embedding_length", 4096)
+ .uint32("llama.attention.head_count", 32)
+ .build()
+
+ val header = read(bytes)
+
+ assertEquals(32L, header?.headCount)
+ assertNull(header?.headCountKv)
+ }
+
+ @Test
+ fun givenShapeKeysBehindATokenizerArray_whenRead_thenTheArrayIsSkippedAndTheKeysAreFound() {
+ val bytes = GgufWriter()
+ .stringArray("tokenizer.ggml.tokens", List(500) { "token$it" })
+ .string("general.architecture", "qwen2")
+ .uint32("qwen2.block_count", 28)
+ .uint32("qwen2.embedding_length", 1536)
+ .uint32("qwen2.attention.head_count", 12)
+ .uint32("qwen2.attention.head_count_kv", 2)
+ .build()
+
+ val header = read(bytes)
+
+ assertEquals(28L, header?.blockCount)
+ assertEquals(2L, header?.headCountKv)
+ }
+
+ @Test
+ fun givenSixtyFourBitShapeValues_whenRead_thenTheyAreReadAtTheRightWidth() {
+ val bytes = GgufWriter()
+ .string("general.architecture", "llama")
+ .uint64("llama.block_count", 32)
+ .uint64("llama.embedding_length", 4096)
+ .uint64("llama.attention.head_count", 32)
+ .build()
+
+ val header = read(bytes)
+
+ assertEquals(32L, header?.blockCount)
+ assertEquals(4096L, header?.embeddingLength)
+ }
+
+ @Test
+ fun givenAShapeValueOfAnUnexpectedType_whenRead_thenOnlyThatValueIsLostAndTheRestSurvive() {
+ // The value still has to be consumed, or every later key would be read from its middle.
+ val bytes = GgufWriter()
+ .string("general.architecture", "llama")
+ .float32("llama.block_count", 32f)
+ .uint32("llama.embedding_length", 4096)
+ .uint32("llama.attention.head_count", 32)
+ .build()
+
+ val header = read(bytes)
+
+ assertNull(header?.blockCount)
+ assertEquals(4096L, header?.embeddingLength)
+ assertEquals(32L, header?.headCount)
+ }
+
+ @Test
+ fun givenAShapeValueOfAnAbsurdMagnitude_whenRead_thenItIsReturnedUnchangedForTheEstimatorToReject() {
+ // The reader reports what the file declares; the plausibility ceiling lives in the estimator.
+ val bytes = GgufWriter()
+ .string("general.architecture", "llama")
+ .uint64("llama.block_count", 1L shl 61)
+ .uint32("llama.attention.head_count", 32)
+ .build()
+
+ val header = read(bytes)
+
+ assertEquals(1L shl 61, header?.blockCount)
+ assertEquals(32L, header?.headCount)
+ }
+
+ @Test
+ fun givenAFileThatIsNotGguf_whenRead_thenThereIsNoHeader() {
+ assertNull(read("This is a text file, not a model".toByteArray()))
+ }
+
+ @Test
+ fun givenATruncatedHeader_whenRead_thenThereIsNoHeader() {
+ val truncated = gemmaHeader().copyOfRange(0, 40)
+
+ assertNull(read(truncated))
+ }
+
+ @Test
+ fun givenAnUnopenableFile_whenRead_thenThereIsNoHeader() {
+ assertNull(GgufHeaderReader.read { null })
+ }
+
+ @Test
+ fun givenAStreamThatFailsMidRead_whenRead_thenThereIsNoHeader() {
+ val failing = object : InputStream() {
+ override fun read(): Int = throw IOException("device detached")
+ }
+
+ assertNull(GgufHeaderReader.read { failing })
+ }
+
+ @Test
+ fun givenAHeaderClaimingAnAbsurdEntryCount_whenRead_thenThereIsNoHeader() {
+ // Truncating is worse than giving up: a head_count_kv never reached reads as plain MHA.
+ val bytes = GgufWriter()
+ .string("general.architecture", "llama")
+ .build()
+ .let { corruptEntryCount(it) }
+
+ assertNull(read(bytes))
+ }
+
+ @Test(timeout = 30_000)
+ fun givenAnArrayCountThatWouldRunOnPastTheMetadata_whenRead_thenTheParseIsAbandoned() {
+ // An array count is unbounded, so without a cap this walks the whole multi-GB file.
+ val bytes = GgufWriter()
+ .lyingStringArray("tokenizer.ggml.tokens", declaredCount = Long.MAX_VALUE / 2, values = listOf("a"))
+ .string("general.architecture", "llama")
+ .build()
+
+ assertNull(GgufHeaderReader.read { EndlessStream(bytes) })
+ }
+
+ @Test(timeout = 30_000)
+ fun givenOneStringValueClaimingAnEnormousLength_whenRead_thenTheByteBudgetStillStopsIt() {
+ // One huge skip, not many small ones, so charging on completion never gets the chance.
+ val bytes = GgufWriter()
+ .lyingString("general.description", declaredLength = Long.MAX_VALUE / 2)
+ .string("general.architecture", "llama")
+ .build()
+
+ assertNull(GgufHeaderReader.read { EndlessStream(bytes) })
+ }
+
+ @Test(timeout = 30_000)
+ fun givenAnArrayNestedInsideAnArray_whenRead_thenItIsRejectedRatherThanFollowed() {
+ // Recursing per nesting level ends in a StackOverflowError, not an Exception.
+ val bytes = GgufWriter()
+ .string("general.architecture", "llama")
+ .nestedArray("tokenizer.ggml.merges")
+ .build()
+
+ assertNull(read(bytes))
+ }
+
+ @Test
+ fun givenShapeKeysForAnotherArchitecture_whenRead_thenOnlyThisModelsAreReturned() {
+ // A multimodal file carries the vision tower's shape alongside the language model's.
+ val bytes = GgufWriter()
+ .string("general.architecture", "qwen2")
+ .uint32("clip.vision.block_count", 27)
+ .uint32("clip.vision.embedding_length", 1152)
+ .uint32("clip.vision.attention.head_count", 16)
+ .uint32("qwen2.block_count", 28)
+ .uint32("qwen2.embedding_length", 1536)
+ .uint32("qwen2.attention.head_count", 12)
+ .build()
+
+ val header = read(bytes)
+
+ assertEquals(28L, header?.blockCount)
+ assertEquals(1536L, header?.embeddingLength)
+ assertEquals(12L, header?.headCount)
+ }
+
+ @Test
+ fun givenTheArchitectureDeclaredAfterTheShapeKeys_whenRead_thenTheyAreStillAttributed() {
+ // general.architecture is conventionally first, but the format does not require it.
+ val bytes = GgufWriter()
+ .uint32("llama.block_count", 32)
+ .uint32("llama.attention.head_count", 32)
+ .string("general.architecture", "llama")
+ .build()
+
+ val header = read(bytes)
+
+ assertEquals(32L, header?.blockCount)
+ assertEquals(32L, header?.headCount)
+ }
+
+ @Test
+ fun givenDeclaredKeyAndValueWidths_whenRead_thenTheyAreReturned() {
+ val bytes = GgufWriter()
+ .string("general.architecture", "gemma3")
+ .uint32("gemma3.attention.key_length", 256)
+ .uint32("gemma3.attention.value_length", 256)
+ .build()
+
+ val header = read(bytes)
+
+ assertEquals(256L, header?.keyLength)
+ assertEquals(256L, header?.valueLength)
+ }
+
+ /** Overwrites the entry count (bytes 16..23 of a v3 header) with a huge value. */
+ private fun corruptEntryCount(bytes: ByteArray): ByteArray = bytes.copyOf().also {
+ for (i in 16 until 24) it[i] = 0xFF.toByte()
+ it[23] = 0x00 // keep it positive
+ }
+
+ /**
+ * [prefix], then zeros forever — a stand-in for a multi-gigabyte model whose metadata lies, so
+ * the reader has to stop itself rather than be stopped by end-of-file. The offset is a Long
+ * deliberately: as an Int it wrapped past 2 GB and threw, ending the parse for the wrong reason.
+ */
+ private class EndlessStream(private val prefix: ByteArray) : InputStream() {
+
+ private var offset = 0L
+
+ override fun read(): Int =
+ (if (offset < prefix.size) prefix[offset.toInt()].toInt() and 0xFF else 0)
+ .also { offset++ }
+
+ override fun read(b: ByteArray, off: Int, len: Int): Int {
+ var written = 0
+ while (written < len && offset < prefix.size) {
+ b[off + written] = prefix[offset.toInt()]
+ offset++
+ written++
+ }
+ if (written < len) {
+ b.fill(0, off + written, off + len)
+ offset += len - written
+ }
+ return len
+ }
+ }
+}
diff --git a/ai-assistant/src/test/kotlin/com/itsaky/androidide/plugins/aiassistant/util/GgufWriter.kt b/ai-assistant/src/test/kotlin/com/itsaky/androidide/plugins/aiassistant/util/GgufWriter.kt
new file mode 100644
index 00000000..36998648
--- /dev/null
+++ b/ai-assistant/src/test/kotlin/com/itsaky/androidide/plugins/aiassistant/util/GgufWriter.kt
@@ -0,0 +1,107 @@
+package com.itsaky.androidide.plugins.aiassistant.util
+
+import java.io.ByteArrayOutputStream
+
+/**
+ * Writes GGUF headers for tests, so the reader is exercised against real bytes rather than a stub.
+ *
+ * @param version the GGUF version to declare; 1 uses 32-bit lengths and counts, 2+ use 64-bit
+ */
+internal class GgufWriter(private val version: Int = 3) {
+
+ private companion object {
+ const val T_UINT32 = 4
+ const val T_FLOAT32 = 6
+ const val T_STRING = 8
+ const val T_ARRAY = 9
+ const val T_UINT64 = 10
+ }
+
+ private val wide = version >= 2
+ private val entries = ByteArrayOutputStream()
+ private var entryCount = 0L
+
+ fun string(key: String, value: String): GgufWriter = entry(key, T_STRING) { writeString(value) }
+
+ fun uint32(key: String, value: Long): GgufWriter = entry(key, T_UINT32) { writeU32(value) }
+
+ fun uint64(key: String, value: Long): GgufWriter = entry(key, T_UINT64) { writeU64(value) }
+
+ fun float32(key: String, value: Float): GgufWriter =
+ entry(key, T_FLOAT32) { writeU32(java.lang.Float.floatToIntBits(value).toLong() and 0xFFFFFFFFL) }
+
+ /** A tokenizer-sized value, so tests can check that the reader skips past one correctly. */
+ fun stringArray(key: String, values: List): GgufWriter = entry(key, T_ARRAY) {
+ writeU32(T_STRING.toLong())
+ writeCount(values.size.toLong())
+ values.forEach { writeString(it) }
+ }
+
+ /**
+ * An array whose declared length lies about how much data follows, for testing the reader's
+ * bounds. A real file corrupted mid-download looks like this.
+ *
+ * @param declaredCount the element count written to the file, however untrue
+ * @param values the elements actually written
+ */
+ fun lyingStringArray(key: String, declaredCount: Long, values: List): GgufWriter =
+ entry(key, T_ARRAY) {
+ writeU32(T_STRING.toLong())
+ writeCount(declaredCount)
+ values.forEach { writeString(it) }
+ }
+
+ /**
+ * A string value whose declared length lies about how much data follows. Distinct from
+ * [lyingStringArray]: this is one enormous skip rather than many small ones, which is what a
+ * byte budget charged only on completion fails to catch.
+ */
+ fun lyingString(key: String, declaredLength: Long): GgufWriter =
+ entry(key, T_STRING) { writeCount(declaredLength) }
+
+ /**
+ * An array whose elements are themselves arrays. The format forbids this, and following one
+ * recurses a level per nesting — the shape that used to end in a StackOverflowError.
+ */
+ fun nestedArray(key: String, declaredCount: Long = 1L): GgufWriter = entry(key, T_ARRAY) {
+ writeU32(T_ARRAY.toLong())
+ writeCount(declaredCount)
+ }
+
+ /** @return the complete header: magic, version, tensor count, entry count, then the entries. */
+ fun build(): ByteArray {
+ val out = ByteArrayOutputStream()
+ out.write("GGUF".toByteArray(Charsets.US_ASCII))
+ out.writeU32(version.toLong())
+ out.writeCount(0L) // tensor count
+ out.writeCount(entryCount)
+ out.write(entries.toByteArray())
+ return out.toByteArray()
+ }
+
+ private fun entry(key: String, type: Int, writeValue: ByteArrayOutputStream.() -> Unit): GgufWriter {
+ entries.writeString(key)
+ entries.writeU32(type.toLong())
+ entries.writeValue()
+ entryCount++
+ return this
+ }
+
+ private fun ByteArrayOutputStream.writeString(value: String) {
+ val bytes = value.toByteArray(Charsets.UTF_8)
+ writeCount(bytes.size.toLong())
+ write(bytes)
+ }
+
+ private fun ByteArrayOutputStream.writeCount(value: Long) {
+ if (wide) writeU64(value) else writeU32(value)
+ }
+
+ private fun ByteArrayOutputStream.writeU32(value: Long) {
+ for (shift in 0 until 4) write(((value shr (8 * shift)) and 0xFF).toInt())
+ }
+
+ private fun ByteArrayOutputStream.writeU64(value: Long) {
+ for (shift in 0 until 8) write(((value shr (8 * shift)) and 0xFF).toInt())
+ }
+}
diff --git a/ai-assistant/src/test/kotlin/com/itsaky/androidide/plugins/aiassistant/viewmodel/AiSettingsViewModelMemoryTest.kt b/ai-assistant/src/test/kotlin/com/itsaky/androidide/plugins/aiassistant/viewmodel/AiSettingsViewModelMemoryTest.kt
new file mode 100644
index 00000000..265c7503
--- /dev/null
+++ b/ai-assistant/src/test/kotlin/com/itsaky/androidide/plugins/aiassistant/viewmodel/AiSettingsViewModelMemoryTest.kt
@@ -0,0 +1,290 @@
+package com.itsaky.androidide.plugins.aiassistant.viewmodel
+
+import android.content.ContentResolver
+import android.content.Context
+import android.net.Uri
+import androidx.arch.core.executor.testing.InstantTaskExecutorRule
+import com.itsaky.androidide.plugins.aiassistant.gemini.GeminiCatalogGateway
+import com.itsaky.androidide.plugins.aiassistant.memory.DeviceMemory
+import com.itsaky.androidide.plugins.aiassistant.memory.ModelMemoryGate
+import com.itsaky.androidide.plugins.aiassistant.util.GgufWriter
+import com.itsaky.androidide.plugins.aiassistant.util.ModelFileInfo
+import com.itsaky.androidide.plugins.aiassistant.util.ModelFileSource
+import io.mockk.every
+import io.mockk.mockk
+import io.mockk.mockkStatic
+import io.mockk.unmockkAll
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.test.UnconfinedTestDispatcher
+import kotlinx.coroutines.test.resetMain
+import kotlinx.coroutines.test.runCurrent
+import kotlinx.coroutines.test.runTest
+import kotlinx.coroutines.test.setMain
+import org.junit.After
+import org.junit.Assert.assertEquals
+import org.junit.Assert.assertNull
+import org.junit.Assert.assertTrue
+import org.junit.Before
+import org.junit.Rule
+import org.junit.Test
+import org.junit.rules.TemporaryFolder
+import java.io.File
+import java.io.InputStream
+
+/**
+ * Tests the memory pre-flight end to end through the ViewModel: what the user is asked and what is
+ * persisted — persisting the path is what makes ai-core load, so "not persisted" asserts no load.
+ * Free RAM and the file lookup are faked; the model is a real file with a real GGUF header.
+ */
+@OptIn(ExperimentalCoroutinesApi::class)
+class AiSettingsViewModelMemoryTest {
+
+ /** The ViewModel touches LiveData in its init block and posts to it from the load. */
+ @get:Rule
+ val instantTaskExecutorRule = InstantTaskExecutorRule()
+
+ @get:Rule
+ val tempFolder = TemporaryFolder()
+
+ private val megabyte = 1024L * 1024
+ private val gigabyte = 1024 * megabyte
+
+ /**
+ * gemma-3-1b's shape: a 122,683,392-byte KV cache at full context, plus the 256 MB compute
+ * buffer. A literal, so this asserts the arithmetic rather than restating it.
+ */
+ private val expectedRunBytes = 122_683_392L + 256 * megabyte
+
+ private val dispatcher = UnconfinedTestDispatcher()
+ private lateinit var modelFile: File
+ private lateinit var fileSource: FakeModelFileSource
+
+ @Before
+ fun setUp() {
+ Dispatchers.setMain(dispatcher)
+ modelFile = tempModel("model.gguf", gguf())
+ fileSource = FakeModelFileSource()
+ // Only for GgufFileInspector's pre-check; everything else goes through the file source.
+ mockkStatic(Uri::class)
+ every { Uri.parse(any()) } returns mockk(relaxed = true)
+ }
+
+ @After
+ fun tearDown() {
+ Dispatchers.resetMain()
+ unmockkAll()
+ }
+
+ @Test
+ fun givenPlentyOfFreeMemory_whenAModelIsSelected_thenItIsAcceptedWithoutAWarning() = runTest {
+ val viewModel = viewModel(availableBytes = 8 * gigabyte)
+ val asked = collectWarnings(viewModel)
+
+ viewModel.loadModelFromUri(modelFile.absolutePath, androidContext(modelFile))
+ runCurrent()
+
+ assertTrue(asked.isEmpty())
+ assertEquals(modelFile.absolutePath, viewModel.savedModelPath.value)
+ assertTrue(viewModel.modelLoadingState.value is ModelLoadingState.Loaded)
+ }
+
+ @Test
+ fun givenTooLittleFreeMemory_whenAModelIsSelected_thenTheUserIsAskedWithTheFigures() = runTest {
+ val viewModel = viewModel(availableBytes = 64 * megabyte)
+ val asked = collectWarnings(viewModel)
+
+ viewModel.loadModelFromUri(modelFile.absolutePath, androidContext(modelFile))
+ runCurrent()
+
+ assertEquals(1, asked.size)
+ val warning = asked.single()
+ assertEquals(modelFile.name, warning.modelName)
+ assertEquals(modelFile.length(), warning.loadBytes)
+ assertEquals(expectedRunBytes, warning.runBytes)
+ assertEquals(64 * megabyte, warning.availableBytes)
+ assertEquals(ModelMemoryGate.Severity.INSUFFICIENT, warning.severity)
+ }
+
+ @Test
+ fun givenTheMemoryWarning_whenTheUserCancels_thenTheModelIsNeverPersisted() = runTest {
+ val viewModel = viewModel(availableBytes = 64 * megabyte)
+ collectWarnings(viewModel)
+ viewModel.loadModelFromUri(modelFile.absolutePath, androidContext(modelFile))
+ runCurrent()
+
+ viewModel.onMemoryWarningDecision(proceed = false)
+ runCurrent()
+
+ assertNull(viewModel.savedModelPath.value)
+ assertEquals(ModelLoadingState.Idle, viewModel.modelLoadingState.value)
+ }
+
+ @Test
+ fun givenTheMemoryWarning_whenTheUserCancels_thenTheDocumentGrantIsGivenBack() = runTest {
+ // The grant is taken before the check runs, and its table has a hard per-app limit.
+ val viewModel = viewModel(availableBytes = 64 * megabyte)
+ collectWarnings(viewModel)
+ viewModel.loadModelFromUri(modelFile.absolutePath, androidContext(modelFile))
+ runCurrent()
+
+ viewModel.onMemoryWarningDecision(proceed = false)
+ runCurrent()
+
+ assertEquals(listOf(modelFile.absolutePath), fileSource.released)
+ }
+
+ @Test
+ fun givenTheMemoryWarning_whenTheUserProceeds_thenTheGrantIsKept() = runTest {
+ val viewModel = viewModel(availableBytes = 64 * megabyte)
+ collectWarnings(viewModel)
+ viewModel.loadModelFromUri(modelFile.absolutePath, androidContext(modelFile))
+ runCurrent()
+
+ viewModel.onMemoryWarningDecision(proceed = true)
+ runCurrent()
+
+ assertEquals(emptyList(), fileSource.released)
+ }
+
+ @Test
+ fun givenTheMemoryWarning_whenTheUserProceeds_thenTheModelIsPersistedAnyway() = runTest {
+ val viewModel = viewModel(availableBytes = 64 * megabyte)
+ collectWarnings(viewModel)
+ viewModel.loadModelFromUri(modelFile.absolutePath, androidContext(modelFile))
+ runCurrent()
+
+ viewModel.onMemoryWarningDecision(proceed = true)
+ runCurrent()
+
+ assertEquals(modelFile.absolutePath, viewModel.savedModelPath.value)
+ assertTrue(viewModel.modelLoadingState.value is ModelLoadingState.Loaded)
+ }
+
+ @Test
+ fun givenUnreadableFreeMemory_whenAModelIsSelected_thenItIsAcceptedRatherThanQuestioned() = runTest {
+ // Failing open: a device we cannot measure must not be told its model won't fit.
+ val viewModel = viewModel(availableBytes = null)
+ val asked = collectWarnings(viewModel)
+
+ viewModel.loadModelFromUri(modelFile.absolutePath, androidContext(modelFile))
+ runCurrent()
+
+ assertTrue(asked.isEmpty())
+ assertEquals(modelFile.absolutePath, viewModel.savedModelPath.value)
+ }
+
+ @Test
+ fun givenAnUnknownFileSize_whenAModelIsSelected_thenItIsAcceptedRatherThanQuestioned() = runTest {
+ fileSource.sizeOverride = null
+ val viewModel = viewModel(availableBytes = 64 * megabyte)
+ val asked = collectWarnings(viewModel)
+
+ viewModel.loadModelFromUri(modelFile.absolutePath, androidContext(modelFile))
+ runCurrent()
+
+ assertTrue(asked.isEmpty())
+ assertEquals(modelFile.absolutePath, viewModel.savedModelPath.value)
+ }
+
+ @Test
+ fun givenAModelWithoutAReadableHeader_whenSelected_thenTheSizeBasedEstimateIsUsed() = runTest {
+ // Valid magic, nothing usable behind it: the estimate falls back instead of vanishing.
+ val headerless = tempModel("headerless.gguf", "GGUF".toByteArray() + ByteArray(8))
+ val viewModel = viewModel(availableBytes = 64 * megabyte)
+ val asked = collectWarnings(viewModel)
+
+ viewModel.loadModelFromUri(headerless.absolutePath, androidContext(headerless))
+ runCurrent()
+
+ assertEquals(256 * megabyte, asked.single().runBytes)
+ }
+
+ @Test
+ fun givenAFileThatIsNotAModel_whenSelected_thenItIsRejectedBeforeTheMemoryCheck() = runTest {
+ val notAModel = tempModel("notes.txt", "nowhere near a model file".toByteArray())
+ val viewModel = viewModel(availableBytes = 64 * megabyte)
+ val asked = collectWarnings(viewModel)
+
+ viewModel.loadModelFromUri(notAModel.absolutePath, androidContext(notAModel))
+ runCurrent()
+
+ assertTrue(asked.isEmpty())
+ assertNull(viewModel.savedModelPath.value)
+ assertTrue(viewModel.modelLoadingState.value is ModelLoadingState.Error)
+ }
+
+ private fun viewModel(availableBytes: Long?) = AiSettingsViewModel(
+ getContext = { null },
+ ioDispatcher = dispatcher,
+ catalogGateway = mockk(relaxed = true),
+ deviceMemory = DeviceMemory { availableBytes },
+ modelFiles = fileSource,
+ )
+
+ /** Collects the warnings the ViewModel raises, so tests can assert on what the user is shown. */
+ private fun kotlinx.coroutines.test.TestScope.collectWarnings(
+ viewModel: AiSettingsViewModel
+ ): List {
+ val asked = mutableListOf()
+ backgroundScope.launch { viewModel.modelMemoryWarnings.collect { asked += it } }
+ return asked
+ }
+
+ /**
+ * Serves the model file over a mocked resolver, which is all GgufFileInspector's magic-byte
+ * pre-check needs. Name, size and the header read go through [fileSource] instead.
+ */
+ private fun androidContext(file: File): Context {
+ val resolver = mockk()
+ every { resolver.openInputStream(any()) } answers { file.inputStream() }
+ return mockk(relaxed = true) {
+ every { contentResolver } returns resolver
+ }
+ }
+
+ private fun tempModel(name: String, bytes: ByteArray): File =
+ tempFolder.newFile(name).apply { writeBytes(bytes) }
+
+ private fun gguf(): ByteArray = GgufWriter()
+ .string("general.architecture", "gemma3")
+ .uint32("gemma3.block_count", 26)
+ .uint32("gemma3.embedding_length", 1152)
+ .uint32("gemma3.attention.head_count", 4)
+ .uint32("gemma3.attention.head_count_kv", 1)
+ .build()
+
+ /**
+ * Reads the real files these tests write, with no Android framework on the path — which is what
+ * the ViewModel taking a [ModelFileSource] instead of a ContentResolver buys.
+ */
+ private class FakeModelFileSource : ModelFileSource {
+
+ /** Set to null to simulate a provider that will not report a size. */
+ var sizeOverride: Long? = UNSET
+
+ val released = mutableListOf()
+
+ override fun info(context: Context, uriString: String): ModelFileInfo {
+ val file = File(uriString)
+ val size = if (sizeOverride == UNSET) file.length().takeIf { it > 0L } else sizeOverride
+ return ModelFileInfo(file.name, size)
+ }
+
+ override fun openStream(context: Context, uriString: String): InputStream? =
+ File(uriString).takeIf { it.isFile }?.inputStream()
+
+ override fun fallbackDisplayName(uriOrPath: String): String =
+ uriOrPath.substringAfterLast('/')
+
+ override fun releaseAccess(context: Context, uriString: String) {
+ released += uriString
+ }
+
+ private companion object {
+ /** Distinguishes "use the real file length" from a deliberate null. */
+ const val UNSET = Long.MIN_VALUE
+ }
+ }
+}
diff --git a/ai-assistant/src/test/kotlin/com/itsaky/androidide/plugins/aiassistant/viewmodel/UserConfirmationTest.kt b/ai-assistant/src/test/kotlin/com/itsaky/androidide/plugins/aiassistant/viewmodel/UserConfirmationTest.kt
new file mode 100644
index 00000000..113c61af
--- /dev/null
+++ b/ai-assistant/src/test/kotlin/com/itsaky/androidide/plugins/aiassistant/viewmodel/UserConfirmationTest.kt
@@ -0,0 +1,177 @@
+package com.itsaky.androidide.plugins.aiassistant.viewmodel
+
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.async
+import kotlinx.coroutines.cancelAndJoin
+import kotlinx.coroutines.flow.first
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.test.runCurrent
+import kotlinx.coroutines.test.runTest
+import org.junit.Assert.assertEquals
+import org.junit.Assert.assertFalse
+import org.junit.Assert.assertTrue
+import org.junit.Test
+
+/**
+ * Tests for the await-the-user primitive behind the memory warning.
+ *
+ * This is where the feature's concurrency lives, so it is tested directly: whether a decision is
+ * delivered, what happens to a question nobody answers, and that two selections cannot interleave.
+ */
+@OptIn(ExperimentalCoroutinesApi::class)
+class UserConfirmationTest {
+
+ @Test
+ fun givenAQuestion_whenTheUserProceeds_thenAskReturnsTrue() = runTest {
+ val confirmation = UserConfirmation()
+
+ val answer = async { confirmation.ask("model.gguf") }
+ assertEquals("model.gguf", confirmation.requests.first())
+ confirmation.answer(true)
+
+ assertTrue(answer.await())
+ }
+
+ @Test
+ fun givenAQuestion_whenTheUserCancels_thenAskReturnsFalse() = runTest {
+ val confirmation = UserConfirmation()
+
+ val answer = async { confirmation.ask("model.gguf") }
+ confirmation.requests.first()
+ confirmation.answer(false)
+
+ assertFalse(answer.await())
+ }
+
+ @Test
+ fun givenAQuestionOnTheWay_whenNobodyHasAnsweredYet_thenAskIsStillWaiting() = runTest {
+ val confirmation = UserConfirmation()
+
+ val answer = async { confirmation.ask("model.gguf") }
+ runCurrent()
+
+ assertTrue(answer.isActive)
+ confirmation.answer(false)
+ assertFalse(answer.await())
+ }
+
+ @Test
+ fun givenAnAnswerWithNothingPending_whenTheNextQuestionIsAsked_thenItIsNotPreAnswered() = runTest {
+ // A stray reply from a dismissed dialog must not decide the following selection.
+ val confirmation = UserConfirmation()
+ confirmation.answer(true)
+
+ val answer = async { confirmation.ask("model.gguf") }
+ runCurrent()
+
+ assertTrue(answer.isActive)
+ confirmation.answer(false)
+ assertFalse(answer.await())
+ }
+
+ @Test
+ fun givenAnAnsweredQuestion_whenAnsweredAgain_thenTheDuplicateIsIgnored() = runTest {
+ // A recreated dialog reporting the same decision twice is harmless.
+ val confirmation = UserConfirmation()
+
+ val answer = async { confirmation.ask("model.gguf") }
+ confirmation.requests.first()
+ confirmation.answer(true)
+ confirmation.answer(false)
+
+ assertTrue(answer.await())
+ }
+
+ @Test
+ fun givenAQuestionAwaitingAnAnswer_whenAnotherIsAsked_thenItWaitsItsTurn() = runTest {
+ val confirmation = UserConfirmation()
+ val asked = mutableListOf()
+ backgroundScope.launch { confirmation.requests.collect { asked += it } }
+
+ val first = async { confirmation.ask("first.gguf") }
+ val second = async { confirmation.ask("second.gguf") }
+ runCurrent()
+
+ // Only one question is outstanding, so an answer can never be attributed to the wrong one.
+ assertEquals(listOf("first.gguf"), asked)
+
+ confirmation.answer(true)
+ runCurrent()
+ assertEquals(listOf("first.gguf", "second.gguf"), asked)
+
+ confirmation.answer(false)
+ assertTrue(first.await())
+ assertFalse(second.await())
+ }
+
+ @Test
+ fun givenAQuestionRaised_whenTheCollectorIsReplaced_thenTheNewOneIsStillAsked() = runTest {
+ // A rotation replaces the collector; the question must not vanish with the old view.
+ val confirmation = UserConfirmation()
+ val firstCollector = launch { confirmation.requests.collect { } }
+
+ val answer = async { confirmation.ask("model.gguf") }
+ runCurrent()
+ firstCollector.cancelAndJoin()
+
+ assertEquals("model.gguf", confirmation.requests.first())
+ confirmation.answer(true)
+ assertTrue(answer.await())
+ }
+
+ @Test
+ fun givenAnAnsweredQuestion_whenACollectorSubscribesAfterwards_thenItIsNotAskedAgain() = runTest {
+ // The withdrawn question must not reappear and re-show the dialog on the next resume.
+ val confirmation = UserConfirmation()
+ val asked = mutableListOf()
+
+ val answer = async { confirmation.ask("model.gguf") }
+ confirmation.requests.first()
+ confirmation.answer(true)
+ assertTrue(answer.await())
+
+ backgroundScope.launch { confirmation.requests.collect { asked += it } }
+ runCurrent()
+
+ assertEquals(emptyList(), asked)
+ }
+
+ @Test
+ fun givenNoQuestionEverAsked_whenInspected_thenNothingIsOutstanding() {
+ // After process death the UI may hold a dialog for a question this object never saw.
+ assertFalse(UserConfirmation().hasOutstandingRequest)
+ }
+
+ @Test
+ fun givenAQuestionOnTheWay_whenInspected_thenItIsOutstandingUntilAnswered() = runTest {
+ val confirmation = UserConfirmation()
+
+ val answer = async { confirmation.ask("model.gguf") }
+ runCurrent()
+ assertTrue(confirmation.hasOutstandingRequest)
+
+ confirmation.answer(true)
+ answer.await()
+
+ assertFalse(confirmation.hasOutstandingRequest)
+ }
+
+ @Test
+ fun givenAnAbandonedQuestion_whenTheCallerIsCancelled_thenTheNextQuestionStillGetsThrough() = runTest {
+ // The ViewModel being cleared mid-dialog must not wedge the next selection.
+ val confirmation = UserConfirmation()
+ val asked = mutableListOf()
+ backgroundScope.launch { confirmation.requests.collect { asked += it } }
+
+ val abandoned = launch { confirmation.ask("abandoned.gguf") }
+ runCurrent()
+ abandoned.cancelAndJoin()
+
+ val answer = async { confirmation.ask("next.gguf") }
+ runCurrent()
+ confirmation.answer(true)
+
+ assertTrue(answer.await())
+ assertEquals(listOf("abandoned.gguf", "next.gguf"), asked)
+ }
+}
diff --git a/ai-core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/LocalLlmBackend.kt b/ai-core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/LocalLlmBackend.kt
index 459bf513..41f192ef 100644
--- a/ai-core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/LocalLlmBackend.kt
+++ b/ai-core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/LocalLlmBackend.kt
@@ -254,6 +254,11 @@ class LocalLlmBackend(private val context: PluginContext) : LlmBackend, Cancella
currentModelPath = null
}
+ // Measured after the unload: availMem excludes the context and batch it just released.
+ ModelLoadDiagnostics.refuseBeforeLoad(availableMemoryBytes())?.let { shortfall ->
+ throw ModelLoadException(loadMessages.describe(shortfall), shortfall)
+ }
+
context.logger.info("Loading model: $resolvedPath")
try {
llama.load(resolvedPath)
diff --git a/ai-core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/ModelLoadDiagnostics.kt b/ai-core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/ModelLoadDiagnostics.kt
index 881112d5..5b92fe94 100644
--- a/ai-core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/ModelLoadDiagnostics.kt
+++ b/ai-core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/ModelLoadDiagnostics.kt
@@ -13,6 +13,12 @@ object ModelLoadDiagnostics {
/** Headroom floor: a small model still needs a KV cache, which scales with context, not size. */
private const val MIN_HEADROOM_BYTES = 256L * 1024 * 1024
+ /**
+ * Compute buffers every load allocates outright, whatever the model's shape. The floor for
+ * [refuseBeforeLoad], kept equal to ai-assistant's `COMPUTE_BUFFER_BYTES` allowance.
+ */
+ private const val MIN_RUN_BYTES = 256L * 1024 * 1024
+
/** Most likely cause of a load failure; the caller resolves each case to a user-facing string. */
sealed interface Diagnosis {
data object FileMissing : Diagnosis
@@ -69,6 +75,22 @@ object ModelLoadDiagnostics {
else Diagnosis.UnsupportedOrCorrupt
}
+ /**
+ * Whether to refuse a load outright, before ggml aborts the process trying it. Weighs only the
+ * compute buffers, so it stays far more permissive than [diagnose]'s attribution headroom:
+ * ai-assistant's pre-flight lets the user proceed, and a refusal here must not overrule that.
+ *
+ * @param availableMemoryBytes free RAM reported by the OS, or negative if unknown
+ * @return the shortfall to refuse with, or null to attempt the load
+ */
+ fun refuseBeforeLoad(availableMemoryBytes: Long): Diagnosis.LowMemory? =
+ // Only a NEGATIVE reading means "unknown"; 0 is a genuine out-of-memory reading.
+ if (availableMemoryBytes in 0L until MIN_RUN_BYTES) {
+ Diagnosis.LowMemory(MIN_RUN_BYTES, availableMemoryBytes)
+ } else {
+ null
+ }
+
// The markers below mirror the messages thrown by LLamaAndroid.load(); keep them in sync with
// that file. Matching on text is best-effort — an unrecognized message falls back to
// UnsupportedOrCorrupt, which is the safe default for a valid-looking file.
diff --git a/ai-core/src/test/kotlin/com/itsaky/androidide/plugins/aicore/ModelLoadDiagnosticsTest.kt b/ai-core/src/test/kotlin/com/itsaky/androidide/plugins/aicore/ModelLoadDiagnosticsTest.kt
index 614e27fe..0ad23296 100644
--- a/ai-core/src/test/kotlin/com/itsaky/androidide/plugins/aicore/ModelLoadDiagnosticsTest.kt
+++ b/ai-core/src/test/kotlin/com/itsaky/androidide/plugins/aicore/ModelLoadDiagnosticsTest.kt
@@ -3,6 +3,7 @@ package com.itsaky.androidide.plugins.aicore
import com.itsaky.androidide.plugins.aicore.ModelLoadDiagnostics.Diagnosis
import org.junit.Assert.assertEquals
import org.junit.Assert.assertFalse
+import org.junit.Assert.assertNull
import org.junit.Assert.assertTrue
import org.junit.Test
import java.io.File
@@ -126,6 +127,32 @@ class ModelLoadDiagnosticsTest {
assertEquals(Diagnosis.UnsupportedOrCorrupt, d)
}
+ @Test
+ fun givenMemoryBelowTheComputeBuffers_whenRefuseBeforeLoad_thenRefused() {
+ val refusal = ModelLoadDiagnostics.refuseBeforeLoad(availableMemoryBytes = 128L shl 20)
+ assertEquals(Diagnosis.LowMemory(256L shl 20, 128L shl 20), refusal)
+ }
+
+ @Test
+ fun givenZeroFreeMemory_whenRefuseBeforeLoad_thenRefused() {
+ // 0 free bytes is a genuine out-of-memory reading, not an unreadable one.
+ assertEquals(0L, ModelLoadDiagnostics.refuseBeforeLoad(0L)?.availableBytes)
+ }
+
+ @Test
+ fun givenUnknownMemory_whenRefuseBeforeLoad_thenAllowed() {
+ // availMem < 0 (unreadable) must fail open, like the pre-flight it defers to.
+ assertNull(ModelLoadDiagnostics.refuseBeforeLoad(-1L))
+ }
+
+ @Test
+ fun givenMemoryTheAiAssistantPreflightCallsTight_whenRefuseBeforeLoad_thenAllowed() {
+ // The gate must never refuse what "Proceed anyway" authorized: an 8.5 GB model with
+ // 1.5 GB free is TIGHT there (free RAM clears its 768 MB run cost), so the load proceeds
+ // even though diagnose()'s size/4 attribution headroom would be 2.1 GB. See ADFA-1798.
+ assertNull(ModelLoadDiagnostics.refuseBeforeLoad(availableMemoryBytes = 1536L shl 20))
+ }
+
@Test
fun givenGgufMagic_whenIsGguf_thenTrue() {
assertTrue(GgufModelInspector.isGguf(tempFile(64).absolutePath))