From 3c1b4de09fd74fd843a8f2a804e69001fd60c6e6 Mon Sep 17 00:00:00 2001
From: John Trujillo
Date: Wed, 5 Aug 2026 14:02:26 -0500
Subject: [PATCH] feat(ai-assistant): warn before loading an oversized local
model
Selecting a .gguf now runs a pre-flight check before the path is persisted, which is what makes ai-core load it. The model's size and its declared shape are weighed against free RAM; if it looks too large the user is asked whether to go ahead, with Cancel as the default and the figures behind the judgement on screen. Declining leaves the previously selected model untouched.
The check fails open throughout: an unreadable size, header or memory reading means no warning rather than a wrong one.
- ModelMemoryEstimator splits mmap'd weights from the KV cache and compute buffers, sized from the model's own shape when the header provides it.
- ModelMemoryGate decides Safe / Risky(TIGHT|INSUFFICIENT) / Unknown.
- GgufHeaderReader reads only the metadata block, bounded by a byte budget, an entry cap and an array cap, and attributes shape keys to their own architecture so a multimodal file's clip.* values stay out.
- UserConfirmation lets the load coroutine await the user's answer, so the flow stays one readable sequence across a rotation.
ADFA-1798
---
ai-assistant/src/main/assets/docs/index.html | 22 ++
.../plugins/aiassistant/AiAssistantPlugin.kt | 37 +++
.../fragments/AiSettingsFragment.kt | 56 +++-
.../fragments/MemoryWarningDialogFragment.kt | 153 +++++++++
.../aiassistant/memory/DeviceMemory.kt | 43 +++
.../memory/ModelMemoryEstimator.kt | 108 +++++++
.../aiassistant/memory/ModelMemoryGate.kt | 53 +++
.../plugins/aiassistant/util/ByteSize.kt | 25 ++
.../aiassistant/util/GgufHeaderReader.kt | 305 ++++++++++++++++++
.../aiassistant/util/ModelFileSource.kt | 132 ++++++++
.../viewmodel/AiSettingsViewModel.kt | 205 +++++++++---
.../aiassistant/viewmodel/UserConfirmation.kt | 68 ++++
ai-assistant/src/main/res/values/strings.xml | 8 +
.../memory/ModelMemoryEstimatorTest.kt | 180 +++++++++++
.../aiassistant/memory/ModelMemoryGateTest.kt | 79 +++++
.../aiassistant/util/GgufHeaderReaderTest.kt | 283 ++++++++++++++++
.../plugins/aiassistant/util/GgufWriter.kt | 107 ++++++
.../AiSettingsViewModelMemoryTest.kt | 290 +++++++++++++++++
.../viewmodel/UserConfirmationTest.kt | 177 ++++++++++
.../plugins/aicore/LocalLlmBackend.kt | 5 +
.../plugins/aicore/ModelLoadDiagnostics.kt | 22 ++
.../aicore/ModelLoadDiagnosticsTest.kt | 27 ++
22 files changed, 2342 insertions(+), 43 deletions(-)
create mode 100644 ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/fragments/MemoryWarningDialogFragment.kt
create mode 100644 ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/memory/DeviceMemory.kt
create mode 100644 ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/memory/ModelMemoryEstimator.kt
create mode 100644 ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/memory/ModelMemoryGate.kt
create mode 100644 ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/util/ByteSize.kt
create mode 100644 ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/util/GgufHeaderReader.kt
create mode 100644 ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/util/ModelFileSource.kt
create mode 100644 ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/viewmodel/UserConfirmation.kt
create mode 100644 ai-assistant/src/test/kotlin/com/itsaky/androidide/plugins/aiassistant/memory/ModelMemoryEstimatorTest.kt
create mode 100644 ai-assistant/src/test/kotlin/com/itsaky/androidide/plugins/aiassistant/memory/ModelMemoryGateTest.kt
create mode 100644 ai-assistant/src/test/kotlin/com/itsaky/androidide/plugins/aiassistant/util/GgufHeaderReaderTest.kt
create mode 100644 ai-assistant/src/test/kotlin/com/itsaky/androidide/plugins/aiassistant/util/GgufWriter.kt
create mode 100644 ai-assistant/src/test/kotlin/com/itsaky/androidide/plugins/aiassistant/viewmodel/AiSettingsViewModelMemoryTest.kt
create mode 100644 ai-assistant/src/test/kotlin/com/itsaky/androidide/plugins/aiassistant/viewmodel/UserConfirmationTest.kt
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))