fix(composer): address the audit of the restored AI writing helper

Bugs - Emptying the composer left the proposals on screen. precomputeAiResults early-returned on short text without clearing state, and cancel() — which runs after a post is sent — resets every other suggestion source but not this one. The freshly emptied composer kept showing proposals for the note just published, and "Use This" pasted it back in. - The client caches were plain HashMaps written by the nine tone coroutines at once. Two tones mapped to the same rewriter, so every batch raced on the same key and orphaned a Rewriter nothing would close. They are now ConcurrentHashMaps built through computeIfAbsent, and close() drains them. - The assistant held the Activity context inside a ViewModel that outlives it. It now keeps the application context, and the screen passes that too, as MLKitImageLabelService already does. - DOWNLOADABLE was folded into "unavailable" and nothing ever called downloadFeature(), so on a device whose model had not been fetched the feature could never start. Status is now re-read (throttled) while it is not ready, and the model is requested once when the user has the setting on. - lastComputedText was stamped before inference, so a cancelled run marked that text as done and returning to it showed nothing. It is stamped after the run completes. - The Settings toggle was read as a plain StateFlow value, so turning it off did not hide the panel. The screen collects it now. - precomputeAiResults/showAiPanel touch a lateinit accountViewModel; they now guard it like the functions above them. Performance - Language detection ran once per tone over identical text; it is memoized per text, so a batch detects once instead of nine times. - MORE_DIRECT and PUNCHY issued the same request as PROFESSIONAL and SHORTER — ML Kit has no other output type for them — so two of nine inferences were wasted and two chip pairs rendered identical text. Both tones are dropped. - Applying a proposal cleared lastComputedText, and the programmatic edit re-entered onMessageChanged, so accepting a suggestion immediately queued a fresh batch over it. It now remembers the applied text. - Inference blocked on future.get(), which coroutine cancellation cannot stop, so abandoned batches kept running. Futures are awaited through suspendCancellableCoroutine and cancelled with the coroutine. - Drafts under 20 characters no longer spend the model at all, and proposals identical to the draft are dropped instead of becoming a chip. Cleanup - Deletes MockWritingAssistant (shipped in main behind a dead flag, carrying its own "remove before shipping" note) and the unused AiWritingHelpButton. - Hides the Settings tile on F-Droid, where the assistant is a no-op. - Panel takes an ImmutableMap; the ML Kit language constants are mapped explicitly instead of relying on the two APIs numbering them alike. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01YRfjXrjHdJ435kQwyp2HnL
2026-10-06 11:48:24 +00:00 · 2026-08-29 20:18:34 +00:00
parent 3eff5f2d87
commit 6f97faf167
12 changed files with 372 additions and 280 deletions
@@ -23,6 +23,8 @@ package com.vitorpamplona.amethyst.service.ai
 class NoOpWritingAssistant : WritingAssistant {
    override suspend fun checkAvailability(): WritingAssistantStatus = WritingAssistantStatus.Unavailable

+    override suspend fun requestDownload(): WritingAssistantStatus = WritingAssistantStatus.Unavailable
+
    override suspend fun transform(
        text: String,
        tone: WritingTone,
@@ -23,6 +23,9 @@ package com.vitorpamplona.amethyst.service.ai
 import android.content.Context

 object WritingAssistantFactory {
+    /** Whether this flavor ships a real assistant. Drives the Settings tile. */
+    const val IS_SUPPORTED = false
+
    @Suppress("UNUSED_PARAMETER")
    fun create(context: Context): WritingAssistant = NoOpWritingAssistant()
 }
@@ -1,99 +0,0 @@
-/*
- * Copyright (c) 2025 Vitor Pamplona
- *
- * Permission is hereby granted, free of charge, to any person obtaining a copy of
- * this software and associated documentation files (the "Software"), to deal in
- * the Software without restriction, including without limitation the rights to use,
- * copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the
- * Software, and to permit persons to whom the Software is furnished to do so,
- * subject to the following conditions:
- *
- * The above copyright notice and this permission notice shall be included in all
- * copies or substantial portions of the Software.
- *
- * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
- * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
- * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
- * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN
- * AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
- * WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
- */
-package com.vitorpamplona.amethyst.service.ai
-
-import kotlinx.coroutines.delay
-
-/**
- * TODO: Remove before shipping. Debug-only mock for testing the AI writing help UI
- * on devices without Gemini Nano / AICore support.
- */
-class MockWritingAssistant : WritingAssistant {
-    override suspend fun checkAvailability(): WritingAssistantStatus = WritingAssistantStatus.Available
-
-    override suspend fun transform(
-        text: String,
-        tone: WritingTone,
-    ): WritingResult {
-        delay(800)
-
-        val transformed =
-            when (tone) {
-                WritingTone.CORRECT -> {
-                    correctMock(text)
-                }
-
-                WritingTone.REPHRASE -> {
-                    "Here's another way to put it: $text"
-                }
-
-                WritingTone.SHORTER -> {
-                    text
-                        .split(".")
-                        .firstOrNull()
-                        ?.trim()
-                        ?.plus(".") ?: text
-                }
-
-                WritingTone.ELABORATE -> {
-                    "$text Furthermore, this point deserves deeper consideration and nuance."
-                }
-
-                WritingTone.FRIENDLY -> {
-                    "Hey! $text Hope that makes sense! :)"
-                }
-
-                WritingTone.PROFESSIONAL -> {
-                    "I would like to bring to your attention the following: $text"
-                }
-
-                WritingTone.MORE_DIRECT -> {
-                    text.replace("I think ", "").replace("maybe ", "").replace("perhaps ", "")
-                }
-
-                WritingTone.PUNCHY -> {
-                    text.uppercase().replace(".", "!")
-                }
-
-                WritingTone.EMOJIFY -> {
-                    "$text \uD83D\uDE80\uD83D\uDD25\u2728"
-                }
-            }
-
-        return WritingResult(
-            originalText = text,
-            transformedText = transformed,
-            tone = tone,
-        )
-    }
-
-    private fun correctMock(text: String): String =
-        text
-            .replace("teh ", "the ")
-            .replace("dont ", "don't ")
-            .replace("cant ", "can't ")
-            .replace("wont ", "won't ")
-            .replace("i ", "I ")
-
-    override fun close() {
-        // no-op: mock holds no native handles or background workers to release.
-    }
-}
@@ -25,6 +25,15 @@ import androidx.compose.runtime.Immutable
 interface WritingAssistant {
    suspend fun checkAvailability(): WritingAssistantStatus

+    /**
+     * Asks the platform to fetch the on-device model when [checkAvailability] reported
+     * [WritingAssistantStatus.Downloadable]. Returns the status after the attempt.
+     *
+     * Implementations must be safe to call repeatedly: only the first call per instance
+     * starts a download, later ones just report the current status.
+     */
+    suspend fun requestDownload(): WritingAssistantStatus
+
    suspend fun transform(
        text: String,
        tone: WritingTone,
@@ -40,8 +49,6 @@ enum class WritingTone {
    ELABORATE,
    FRIENDLY,
    PROFESSIONAL,
-    MORE_DIRECT,
-    PUNCHY,
    EMOJIFY,
 }

@@ -50,6 +57,9 @@ sealed class WritingAssistantStatus {

    data object Unavailable : WritingAssistantStatus()

+    /** The device supports the model but it has not been fetched yet. */
+    data object Downloadable : WritingAssistantStatus()
+
    data object Downloading : WritingAssistantStatus()
 }

@@ -1,54 +0,0 @@
-/*
- * Copyright (c) 2025 Vitor Pamplona
- *
- * Permission is hereby granted, free of charge, to any person obtaining a copy of
- * this software and associated documentation files (the "Software"), to deal in
- * the Software without restriction, including without limitation the rights to use,
- * copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the
- * Software, and to permit persons to whom the Software is furnished to do so,
- * subject to the following conditions:
- *
- * The above copyright notice and this permission notice shall be included in all
- * copies or substantial portions of the Software.
- *
- * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
- * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
- * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
- * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN
- * AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
- * WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
- */
-package com.vitorpamplona.amethyst.ui.note.creators.aihelp
-
-import androidx.compose.foundation.layout.height
-import androidx.compose.material3.IconButton
-import androidx.compose.material3.MaterialTheme
-import androidx.compose.runtime.Composable
-import androidx.compose.ui.Modifier
-import androidx.compose.ui.unit.dp
-import com.vitorpamplona.amethyst.R
-import com.vitorpamplona.amethyst.commons.icons.symbols.Icon
-import com.vitorpamplona.amethyst.commons.icons.symbols.MaterialSymbols
-import com.vitorpamplona.amethyst.ui.stringRes
-
-@Composable
-fun AiWritingHelpButton(
-    isActive: Boolean,
-    onClick: () -> Unit,
-) {
-    IconButton(
-        onClick = { onClick() },
-    ) {
-        Icon(
-            symbol = MaterialSymbols.AutoAwesome,
-            contentDescription = stringRes(R.string.ai_writing_help),
-            modifier = Modifier.height(22.dp),
-            tint =
-                if (isActive) {
-                    MaterialTheme.colorScheme.primary
-                } else {
-                    MaterialTheme.colorScheme.onBackground
-                },
-        )
-    }
-}
@@ -47,11 +47,12 @@ import com.vitorpamplona.amethyst.R
 import com.vitorpamplona.amethyst.service.ai.WritingResult
 import com.vitorpamplona.amethyst.service.ai.WritingTone
 import com.vitorpamplona.amethyst.ui.stringRes
+import kotlinx.collections.immutable.ImmutableMap

@Composable
 fun AiWritingHelpPanel(
    isVisible: Boolean,
-    readyResults: Map<WritingTone, WritingResult>,
+    readyResults: ImmutableMap<WritingTone, WritingResult>,
    selectedResult: WritingResult?,
    onToneSelected: (WritingTone) -> Unit,
    onApply: () -> Unit,
@@ -163,7 +164,5 @@ private fun toneDisplayName(tone: WritingTone): String =
        WritingTone.ELABORATE -> stringRes(R.string.ai_tone_elaborate)
        WritingTone.FRIENDLY -> stringRes(R.string.ai_tone_friendly)
        WritingTone.PROFESSIONAL -> stringRes(R.string.ai_tone_professional)
-        WritingTone.MORE_DIRECT -> stringRes(R.string.ai_tone_more_direct)
-        WritingTone.PUNCHY -> stringRes(R.string.ai_tone_punchy)
        WritingTone.EMOJIFY -> stringRes(R.string.ai_tone_emojify)
    }
@@ -83,6 +83,7 @@ import com.vitorpamplona.amethyst.R
 import com.vitorpamplona.amethyst.commons.icons.symbols.Icon
 import com.vitorpamplona.amethyst.commons.icons.symbols.MaterialSymbols
 import com.vitorpamplona.amethyst.commons.nip30CustomEmojis.ui.ShowEmojiSuggestionList
+import com.vitorpamplona.amethyst.model.BooleanType
 import com.vitorpamplona.amethyst.ui.actions.StrippingFailureDialog
 import com.vitorpamplona.amethyst.ui.actions.mediaServers.FileServerSelectionRow
 import com.vitorpamplona.amethyst.ui.actions.uploads.MAX_VOICE_RECORD_SECONDS
@@ -180,8 +181,15 @@ fun ShortNotePostScreen(
    val activity = context.getActivity()
    val scope = rememberCoroutineScope()

-    LaunchedEffect(Unit) {
-        postViewModel.initWritingAssistant(context)
+    val proposeAiImprovements by
+        accountViewModel.settings.uiSettingsFlow.automaticallyProposeAiImprovements
+            .collectAsStateWithLifecycle()
+
+    LaunchedEffect(proposeAiImprovements) {
+        if (proposeAiImprovements == BooleanType.ALWAYS) {
+            // The assistant outlives this Activity inside the ViewModel, so it must not hold it.
+            postViewModel.initWritingAssistant(context.applicationContext)
+        }
    }

    LaunchedEffect(postViewModel, accountViewModel) {
@@ -714,8 +722,12 @@ private fun NewPostScreenBody(
            )
        }

+        val proposeAiImprovements by
+            accountViewModel.settings.uiSettingsFlow.automaticallyProposeAiImprovements
+                .collectAsStateWithLifecycle()
+
        AiWritingHelpPanel(
-            isVisible = postViewModel.showAiPanel,
+            isVisible = proposeAiImprovements == BooleanType.ALWAYS && postViewModel.showAiPanel,
            readyResults = postViewModel.aiResults,
            selectedResult = postViewModel.aiSelectedResult,
            onToneSelected = postViewModel::selectAiResult,
@@ -51,7 +51,6 @@ import com.vitorpamplona.amethyst.model.LocalCache
 import com.vitorpamplona.amethyst.model.Note
 import com.vitorpamplona.amethyst.model.User
 import com.vitorpamplona.amethyst.model.accountsCache.AccountCacheState
-import com.vitorpamplona.amethyst.service.ai.MockWritingAssistant
 import com.vitorpamplona.amethyst.service.ai.WritingAssistant
 import com.vitorpamplona.amethyst.service.ai.WritingAssistantFactory
 import com.vitorpamplona.amethyst.service.ai.WritingAssistantStatus
@@ -165,6 +164,9 @@ import com.vitorpamplona.quartz.utils.Log
 import com.vitorpamplona.quartz.utils.RandomInstance
 import com.vitorpamplona.quartz.utils.TimeUtils
 import kotlinx.collections.immutable.ImmutableList
+import kotlinx.collections.immutable.ImmutableMap
+import kotlinx.collections.immutable.persistentMapOf
+import kotlinx.collections.immutable.toImmutableMap
 import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.Job
@@ -184,6 +186,15 @@ enum class UserSuggestionAnchor {
    NOTIFY,
 }

+/** Below this many characters a draft is too short for a rewrite to say anything useful. */
+private const val AI_MIN_TEXT_LENGTH = 20
+
+/** How long typing must pause before spending the device's model on the draft. */
+private const val AI_DEBOUNCE_MS = 1000L
+
+/** Model status crosses a process boundary, so it is only re-read this often. */
+private const val AI_STATUS_RECHECK_SECONDS = 30
+
@Stable
 open class ShortNotePostViewModel :
    ViewModel(),
@@ -423,57 +434,92 @@ open class ShortNotePostViewModel :
        return accountViewModel.account.powDifficultyFor(anticipatedPowKind())
    }

-    // AI Writing Help for testing
-    private val useMockAi = false
+    // --- AI Writing Help -----------------------------------------------------------------

-    var aiResults by mutableStateOf<Map<WritingTone, WritingResult>>(emptyMap())
+    var aiResults by mutableStateOf<ImmutableMap<WritingTone, WritingResult>>(persistentMapOf())
    var aiSelectedResult by mutableStateOf<WritingResult?>(null)
    var aiStatus by mutableStateOf<WritingAssistantStatus>(WritingAssistantStatus.Unavailable)
    private var writingAssistant: WritingAssistant? = null
    private var aiComputeJob: Job? = null
+    private var aiStatusJob: Job? = null
    private var lastComputedText: String = ""
+    private var lastStatusCheckAt: Long = 0

+    /**
+     * Whether there is anything to show. The user's preference is watched by the screen so
+     * that turning the setting off hides the panel right away.
+     */
    val showAiPanel: Boolean
-        get() {
-            val prefEnabled =
-                accountViewModel.settings.uiSettingsFlow.automaticallyProposeAiImprovements.value ==
-                    BooleanType.ALWAYS
-            return prefEnabled && aiStatus is WritingAssistantStatus.Available && aiResults.isNotEmpty()
-        }
+        get() = aiStatus is WritingAssistantStatus.Available && aiResults.isNotEmpty()
+
+    private fun isAiEnabledInSettings(): Boolean =
+        accountViewModel.settings.uiSettingsFlow.automaticallyProposeAiImprovements.value ==
+            BooleanType.ALWAYS

    fun initWritingAssistant(context: Context) {
-        if (writingAssistant == null) {
-            writingAssistant =
-                if (useMockAi) {
-                    MockWritingAssistant()
-                } else {
-                    WritingAssistantFactory.create(context)
-                }
-            viewModelScope.launch(Dispatchers.IO) {
-                aiStatus = writingAssistant?.checkAvailability() ?: WritingAssistantStatus.Unavailable
+        if (writingAssistant != null) return
+        writingAssistant = WritingAssistantFactory.create(context)
+        refreshAiStatus()
+    }
+
+    /**
+     * Re-reads the model status, and asks for the download when the device can run the model
+     * but has not fetched it yet. Runs outside [aiComputeJob] so a keystroke cannot cancel a
+     * download halfway, and is throttled because it crosses into another process.
+     */
+    private fun refreshAiStatus() {
+        val assistant = writingAssistant ?: return
+        if (aiStatusJob?.isActive == true) return
+
+        val now = TimeUtils.now()
+        if (lastStatusCheckAt > 0 && now - lastStatusCheckAt < AI_STATUS_RECHECK_SECONDS) return
+        lastStatusCheckAt = now
+
+        // Stays on the ViewModel's main dispatcher so every AI field is written from one
+        // thread; the assistant moves its own blocking work to IO.
+        aiStatusJob =
+            viewModelScope.launch {
+                val status = assistant.checkAvailability()
+                aiStatus =
+                    if (status is WritingAssistantStatus.Downloadable && isAiEnabledInSettings()) {
+                        assistant.requestDownload()
+                    } else {
+                        status
+                    }
+
+                // The model may have become usable while the draft sat untouched; the next
+                // keystroke should not be what finally starts the proposals.
+                if (aiStatus is WritingAssistantStatus.Available) precomputeAiResults()
            }
-        }
    }

    fun precomputeAiResults() {
-        val assistant = writingAssistant ?: return
-        if (aiStatus !is WritingAssistantStatus.Available) return
-        val prefEnabled =
-            accountViewModel.settings.uiSettingsFlow.automaticallyProposeAiImprovements.value ==
-                BooleanType.ALWAYS
-        if (!prefEnabled) return
+        if (!::accountViewModel.isInitialized) return

        val text = message.text.toString().trim()
-        if (text.isBlank() || text == lastComputedText) return
+        if (text.length < AI_MIN_TEXT_LENGTH) {
+            // Covers the composer being emptied or reset: the old proposals no longer
+            // describe what is in the field, so they must not stay on screen.
+            resetAiState()
+            return
+        }
+        if (text == lastComputedText) return
+
+        val assistant = writingAssistant ?: return
+        if (!isAiEnabledInSettings()) return
+
+        if (aiStatus !is WritingAssistantStatus.Available) {
+            refreshAiStatus()
+            return
+        }

        aiComputeJob?.cancel()
-        aiResults = emptyMap()
+        aiResults = persistentMapOf()
        aiSelectedResult = null

        aiComputeJob =
-            viewModelScope.launch(Dispatchers.IO) {
-                delay(1000)
-                lastComputedText = text
+            viewModelScope.launch {
+                delay(AI_DEBOUNCE_MS)

                val results =
                    coroutineScope {
@@ -482,16 +528,25 @@ open class ShortNotePostViewModel :
                                async {
                                    try {
                                        assistant.transform(text, tone)
+                                    } catch (e: CancellationException) {
+                                        throw e
                                    } catch (e: Exception) {
-                                        if (e is CancellationException) throw e
+                                        Log.w("ShortNotePostViewModel", "Could not run the $tone rewrite", e)
                                        null
                                    }
                                }
                            }.awaitAll()
                            .filterNotNull()
+                            // A model that hands back the input unchanged has nothing to offer.
+                            .filter { it.transformedText.isNotBlank() && it.transformedText != it.originalText }
                            .associateBy { it.tone }
+                            .toImmutableMap()
                    }
+
                aiResults = results
+                // Only now: a run that was cancelled must not mark this text as done, or
+                // coming back to it would show nothing.
+                lastComputedText = text
            }
    }

@@ -501,10 +556,13 @@ open class ShortNotePostViewModel :

    fun applyAiResult() {
        aiSelectedResult?.let {
-            message.setTextAndPlaceCursorAtEnd(it.transformedText)
+            val applied = it.transformedText
+            message.setTextAndPlaceCursorAtEnd(applied)
            aiSelectedResult = null
-            aiResults = emptyMap()
-            lastComputedText = ""
+            aiResults = persistentMapOf()
+            // The edit above re-enters onMessageChanged. Remember the applied text so it
+            // does not immediately trigger a fresh batch of inferences over it.
+            lastComputedText = applied.trim()
            draftTag.newVersion()
        }
    }
@@ -513,6 +571,14 @@ open class ShortNotePostViewModel :
        aiSelectedResult = null
    }

+    private fun resetAiState() {
+        aiComputeJob?.cancel()
+        aiComputeJob = null
+        aiResults = persistentMapOf()
+        aiSelectedResult = null
+        lastComputedText = ""
+    }
+
    fun lnAddress(): String? = account.userProfile().lnAddress()

    fun hasLnAddress(): Boolean = account.userProfile().lnAddress() != null
@@ -1608,6 +1674,8 @@ open class ShortNotePostViewModel :
        iMetaAttachments.reset()

        emojiSuggestions?.reset()
+
+        resetAiState()
    }

    fun deleteMediaToUpload(selected: SelectedMediaProcessing) {
@@ -52,6 +52,7 @@ import com.vitorpamplona.amethyst.commons.service.pow.PoWEstimator
 import com.vitorpamplona.amethyst.model.AccountPoWPreferences
 import com.vitorpamplona.amethyst.model.BooleanType
 import com.vitorpamplona.amethyst.model.UiSettingsFlow
+import com.vitorpamplona.amethyst.service.ai.WritingAssistantFactory
 import com.vitorpamplona.amethyst.service.pow.deviceHashesPerSecond
 import com.vitorpamplona.amethyst.service.pow.formatApproxDuration
 import com.vitorpamplona.amethyst.ui.navigation.navs.INav
@@ -102,13 +103,16 @@ fun ComposeSettingsContent(
                description = R.string.auto_create_drafts_setting_description,
            )
            SettingsDivider()
-            BooleanSwitchTile(
-                flow = sharedPrefs.automaticallyProposeAiImprovements,
-                icon = MaterialSymbols.AutoAwesome,
-                title = R.string.ai_writing_setting_title,
-                description = R.string.ai_writing_setting_description,
-            )
-            SettingsDivider()
+            // The F-Droid build ships no on-device model, so the toggle would do nothing there.
+            if (WritingAssistantFactory.IS_SUPPORTED) {
+                BooleanSwitchTile(
+                    flow = sharedPrefs.automaticallyProposeAiImprovements,
+                    icon = MaterialSymbols.AutoAwesome,
+                    title = R.string.ai_writing_setting_title,
+                    description = R.string.ai_writing_setting_description,
+                )
+                SettingsDivider()
+            }
            BooleanSwitchTile(
                flow = sharedPrefs.useTrackedBroadcasts,
                icon = MaterialSymbols.CellTower,
@@ -4918,7 +4918,6 @@
    <string name="relay_join_request">Relay join request</string>
    <string name="relay_leave_request">Relay leave request</string>

-    <string name="ai_writing_help">AI Writing Help</string>
    <string name="ai_writing_setting_title">Propose text improvements</string>
    <string name="ai_writing_setting_description">Uses an on-device AI model to propose text corrections and tone changes.</string>
    <string name="tracked_broadcasts_setting_title">Tracked broadcasts</string>
@@ -5023,8 +5022,6 @@
    <string name="ai_tone_elaborate">Elaborate</string>
    <string name="ai_tone_friendly">Friendly</string>
    <string name="ai_tone_professional">Professional</string>
-    <string name="ai_tone_more_direct">More Direct</string>
-    <string name="ai_tone_punchy">Punchy</string>
    <string name="ai_tone_emojify">+ Emoji</string>

    <!-- Emoji packs -->
@@ -22,7 +22,10 @@ package com.vitorpamplona.amethyst.service.ai

 import android.content.Context
 import com.google.android.gms.tasks.Tasks
+import com.google.common.util.concurrent.ListenableFuture
+import com.google.mlkit.genai.common.DownloadCallback
 import com.google.mlkit.genai.common.FeatureStatus
+import com.google.mlkit.genai.common.GenAiException
 import com.google.mlkit.genai.proofreading.Proofreader
 import com.google.mlkit.genai.proofreading.ProofreaderOptions
 import com.google.mlkit.genai.proofreading.Proofreading
@@ -32,62 +35,143 @@ import com.google.mlkit.genai.rewriting.RewriterOptions
 import com.google.mlkit.genai.rewriting.Rewriting
 import com.google.mlkit.genai.rewriting.RewritingRequest
 import com.vitorpamplona.amethyst.service.lang.LanguageTranslatorService
+import com.vitorpamplona.quartz.utils.Log
+import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.suspendCancellableCoroutine
+import kotlinx.coroutines.sync.Mutex
+import kotlinx.coroutines.sync.withLock
 import kotlinx.coroutines.withContext
+import java.util.concurrent.ConcurrentHashMap
+import java.util.concurrent.ExecutionException
+import java.util.concurrent.Executor
+import kotlin.coroutines.resume
+import kotlin.coroutines.resumeWithException

+/**
+ * On-device writing assistance backed by ML Kit GenAI (Gemini Nano through AICore).
+ *
+ * Clients are cached per (output type, language) because building one is expensive, and
+ * the composer asks for every tone at once — so the cache is hit concurrently and must be
+ * a concurrent map.
+ */
 class MLKitWritingAssistant(
-    private val context: Context,
+    context: Context,
 ) : WritingAssistant {
-    private var rewriters = mutableMapOf<Long, Rewriter>()
-    private var proofreaders = mutableMapOf<Int, Proofreader>()
+    private val context = context.applicationContext
+
+    private val rewriters = ConcurrentHashMap<Long, Rewriter>()
+    private val proofreaders = ConcurrentHashMap<Int, Proofreader>()
+
+    @Volatile
+    private var closed = false
+
+    private val downloadMutex = Mutex()
+    private var downloadRequested = false
+
+    private val languageMutex = Mutex()
+    private var lastDetectedText: String? = null
+    private var lastDetectedLanguage: WritingLanguage = WritingLanguage.ENGLISH

    private fun rewriterCacheKey(
        @RewriterOptions.OutputType outputType: Int,
-        @RewriterOptions.Language language: Int,
-    ): Long = (outputType.toLong() shl 32) or language.toLong()
+        language: WritingLanguage,
+    ): Long = (outputType.toLong() shl 32) or language.rewriterCode.toLong()

    private fun getRewriter(
        @RewriterOptions.OutputType outputType: Int,
-        @RewriterOptions.Language language: Int,
-    ): Rewriter =
-        rewriters.getOrPut(rewriterCacheKey(outputType, language)) {
-            Rewriting.getClient(
-                RewriterOptions
-                    .builder(context)
-                    .setOutputType(outputType)
-                    .setLanguage(language)
-                    .build(),
-            )
+        language: WritingLanguage,
+    ): Rewriter {
+        check(!closed) { "MLKitWritingAssistant is closed" }
+        val key = rewriterCacheKey(outputType, language)
+        val rewriter =
+            rewriters.computeIfAbsent(key) {
+                Rewriting.getClient(
+                    RewriterOptions
+                        .builder(context)
+                        .setOutputType(outputType)
+                        .setLanguage(language.rewriterCode)
+                        .build(),
+                )
+            }
+        // close() may have run between the check above and the insert. Undo the insert
+        // rather than leaking a client nothing will ever close.
+        if (closed) {
+            rewriters.remove(key)?.close()
+            error("MLKitWritingAssistant is closed")
        }
+        return rewriter
+    }

-    private fun getProofreader(
-        @ProofreaderOptions.Language language: Int,
-    ): Proofreader =
-        proofreaders.getOrPut(language) {
-            Proofreading.getClient(
-                ProofreaderOptions
-                    .builder(context)
-                    .setInputType(ProofreaderOptions.InputType.KEYBOARD)
-                    .setLanguage(language)
-                    .build(),
-            )
+    private fun getProofreader(language: WritingLanguage): Proofreader {
+        check(!closed) { "MLKitWritingAssistant is closed" }
+        val key = language.proofreaderCode
+        val proofreader =
+            proofreaders.computeIfAbsent(key) {
+                Proofreading.getClient(
+                    ProofreaderOptions
+                        .builder(context)
+                        .setInputType(ProofreaderOptions.InputType.KEYBOARD)
+                        .setLanguage(key)
+                        .build(),
+                )
+            }
+        if (closed) {
+            proofreaders.remove(key)?.close()
+            error("MLKitWritingAssistant is closed")
        }
+        return proofreader
+    }

    override suspend fun checkAvailability(): WritingAssistantStatus =
        withContext(Dispatchers.IO) {
            try {
-                val rewriter = getRewriter(RewriterOptions.OutputType.REPHRASE, RewriterOptions.Language.ENGLISH)
-                val status = rewriter.checkFeatureStatus().get()
-                when (status) {
-                    FeatureStatus.AVAILABLE -> WritingAssistantStatus.Available
-                    FeatureStatus.DOWNLOADING -> WritingAssistantStatus.Downloading
-                    else -> WritingAssistantStatus.Unavailable
-                }
+                statusOf(getRewriter(RewriterOptions.OutputType.REPHRASE, WritingLanguage.ENGLISH))
+            } catch (e: CancellationException) {
+                throw e
            } catch (e: Exception) {
+                Log.w(TAG, "Could not read the writing assistant status", e)
                WritingAssistantStatus.Unavailable
            }
        }

+    override suspend fun requestDownload(): WritingAssistantStatus =
+        withContext(Dispatchers.IO) {
+            try {
+                downloadMutex.withLock {
+                    val rewriter = getRewriter(RewriterOptions.OutputType.REPHRASE, WritingLanguage.ENGLISH)
+                    if (!downloadRequested) {
+                        downloadRequested = true
+                        rewriter.downloadFeature(SilentDownloadCallback).await()
+                        // The proofreader ships as its own feature: fetch it too, or the
+                        // CORRECT tone would stay missing forever. A failure here is not
+                        // fatal — the rewriting tones still work.
+                        try {
+                            getProofreader(WritingLanguage.ENGLISH).downloadFeature(SilentDownloadCallback).await()
+                        } catch (e: CancellationException) {
+                            throw e
+                        } catch (e: Exception) {
+                            Log.w(TAG, "Could not download the proofreading model", e)
+                        }
+                    }
+                    statusOf(rewriter)
+                }
+            } catch (e: CancellationException) {
+                throw e
+            } catch (e: Exception) {
+                Log.w(TAG, "Could not download the writing assistant model", e)
+                WritingAssistantStatus.Unavailable
+            }
+        }
+
+    private suspend fun statusOf(rewriter: Rewriter): WritingAssistantStatus =
+        when (rewriter.checkFeatureStatus().await()) {
+            FeatureStatus.AVAILABLE -> WritingAssistantStatus.Available
+            FeatureStatus.DOWNLOADING -> WritingAssistantStatus.Downloading
+            FeatureStatus.DOWNLOADABLE -> WritingAssistantStatus.Downloadable
+            else -> WritingAssistantStatus.Unavailable
+        }
+
    override suspend fun transform(
        text: String,
        tone: WritingTone,
@@ -102,8 +186,6 @@ class MLKitWritingAssistant(
                WritingTone.FRIENDLY -> rewrite(text, RewriterOptions.OutputType.FRIENDLY, language)
                WritingTone.PROFESSIONAL -> rewrite(text, RewriterOptions.OutputType.PROFESSIONAL, language)
                WritingTone.EMOJIFY -> rewrite(text, RewriterOptions.OutputType.EMOJIFY, language)
-                WritingTone.MORE_DIRECT -> rewrite(text, RewriterOptions.OutputType.PROFESSIONAL, language)
-                WritingTone.PUNCHY -> rewrite(text, RewriterOptions.OutputType.SHORTEN, language)
            }

        return WritingResult(
@@ -113,57 +195,122 @@ class MLKitWritingAssistant(
        )
    }

-    private suspend fun detectLanguage(text: String): Int =
-        withContext(Dispatchers.IO) {
-            try {
-                val langTag = Tasks.await(LanguageTranslatorService.identifyLanguage(text))
-                mapLanguageTag(langTag)
-            } catch (e: Exception) {
-                RewriterOptions.Language.ENGLISH
-            }
+    /**
+     * Every tone of a batch runs over the same text, so the detection result is memoized:
+     * the first caller pays for it and the rest read the cached answer.
+     */
+    private suspend fun detectLanguage(text: String): WritingLanguage =
+        languageMutex.withLock {
+            if (lastDetectedText == text) return@withLock lastDetectedLanguage
+
+            val detected =
+                try {
+                    withContext(Dispatchers.IO) {
+                        WritingLanguage.fromTag(Tasks.await(LanguageTranslatorService.identifyLanguage(text)))
+                    }
+                } catch (e: CancellationException) {
+                    throw e
+                } catch (e: Exception) {
+                    Log.w(TAG, "Could not identify the language of the post", e)
+                    WritingLanguage.ENGLISH
+                }
+
+            lastDetectedText = text
+            lastDetectedLanguage = detected
+            detected
        }

    private suspend fun rewrite(
        text: String,
        @RewriterOptions.OutputType outputType: Int,
-        @RewriterOptions.Language language: Int,
-    ): String =
-        withContext(Dispatchers.IO) {
-            val rewriter = getRewriter(outputType, language)
-            val request = RewritingRequest.builder(text).build()
-            val result = rewriter.runInference(request).get()
-            result.results.firstOrNull()?.text ?: text
-        }
+        language: WritingLanguage,
+    ): String {
+        // Building a client touches disk and another process; awaiting the inference does not.
+        val rewriter = withContext(Dispatchers.IO) { getRewriter(outputType, language) }
+        val result = rewriter.runInference(RewritingRequest.builder(text).build()).await()
+        return result.results.firstOrNull()?.text ?: text
+    }

    private suspend fun proofread(
        text: String,
-        @ProofreaderOptions.Language language: Int,
-    ): String =
-        withContext(Dispatchers.IO) {
-            val proofreader = getProofreader(language)
-            val request = ProofreadingRequest.builder(text).build()
-            val result = proofreader.runInference(request).get()
-            result.results.firstOrNull()?.text ?: text
-        }
+        language: WritingLanguage,
+    ): String {
+        val proofreader = withContext(Dispatchers.IO) { getProofreader(language) }
+        val result = proofreader.runInference(ProofreadingRequest.builder(text).build()).await()
+        return result.results.firstOrNull()?.text ?: text
+    }

    override fun close() {
-        rewriters.values.forEach { it.close() }
-        rewriters.clear()
-        proofreaders.values.forEach { it.close() }
-        proofreaders.clear()
+        closed = true
+        rewriters.keys.toList().forEach { rewriters.remove(it)?.close() }
+        proofreaders.keys.toList().forEach { proofreaders.remove(it)?.close() }
+    }
+
+    /**
+     * Bridges a [ListenableFuture] into a cancellable suspend call: cancelling the caller
+     * cancels the inference instead of leaving it running on a thread nobody waits for.
+     */
+    private suspend fun <T> ListenableFuture<T>.await(): T =
+        suspendCancellableCoroutine { continuation ->
+            addListener(
+                {
+                    try {
+                        continuation.resume(get())
+                    } catch (e: CancellationException) {
+                        continuation.cancel(e)
+                    } catch (e: ExecutionException) {
+                        continuation.resumeWithException(e.cause ?: e)
+                    } catch (e: Exception) {
+                        continuation.resumeWithException(e)
+                    }
+                },
+                DIRECT_EXECUTOR,
+            )
+            continuation.invokeOnCancellation { cancel(true) }
+        }
+
+    /**
+     * The two ML Kit APIs declare their own language constants. They happen to share the
+     * same numbering today; this enum keeps our call sites from depending on that.
+     */
+    enum class WritingLanguage(
+        val rewriterCode: Int,
+        val proofreaderCode: Int,
+    ) {
+        ENGLISH(RewriterOptions.Language.ENGLISH, ProofreaderOptions.Language.ENGLISH),
+        JAPANESE(RewriterOptions.Language.JAPANESE, ProofreaderOptions.Language.JAPANESE),
+        KOREAN(RewriterOptions.Language.KOREAN, ProofreaderOptions.Language.KOREAN),
+        GERMAN(RewriterOptions.Language.GERMAN, ProofreaderOptions.Language.GERMAN),
+        FRENCH(RewriterOptions.Language.FRENCH, ProofreaderOptions.Language.FRENCH),
+        ITALIAN(RewriterOptions.Language.ITALIAN, ProofreaderOptions.Language.ITALIAN),
+        SPANISH(RewriterOptions.Language.SPANISH, ProofreaderOptions.Language.SPANISH),
+        ;
+
+        companion object {
+            fun fromTag(tag: String?): WritingLanguage =
+                when (tag?.lowercase()?.take(2)) {
+                    "en" -> ENGLISH
+                    "ja" -> JAPANESE
+                    "ko" -> KOREAN
+                    "de" -> GERMAN
+                    "fr" -> FRENCH
+                    "it" -> ITALIAN
+                    "es" -> SPANISH
+                    else -> ENGLISH
+                }
+        }
+    }
+
+    private object SilentDownloadCallback : DownloadCallback {
+        override fun onDownloadFailed(e: GenAiException) {
+            Log.w(TAG, "Writing assistant model download failed", e)
+        }
    }

    companion object {
-        fun mapLanguageTag(tag: String?): Int =
-            when (tag?.lowercase()?.take(2)) {
-                "en" -> RewriterOptions.Language.ENGLISH
-                "ja" -> RewriterOptions.Language.JAPANESE
-                "ko" -> RewriterOptions.Language.KOREAN
-                "de" -> RewriterOptions.Language.GERMAN
-                "fr" -> RewriterOptions.Language.FRENCH
-                "it" -> RewriterOptions.Language.ITALIAN
-                "es" -> RewriterOptions.Language.SPANISH
-                else -> RewriterOptions.Language.ENGLISH
-            }
+        private const val TAG = "MLKitWritingAssistant"
+
+        /** Completes the continuation on whichever thread finished the future. */
+        private val DIRECT_EXECUTOR = Executor { it.run() }
    }
 }
@@ -23,5 +23,8 @@ package com.vitorpamplona.amethyst.service.ai
 import android.content.Context

 object WritingAssistantFactory {
+    /** Whether this flavor ships a real assistant. Drives the Settings tile. */
+    const val IS_SUPPORTED = true
+
    fun create(context: Context): WritingAssistant = MLKitWritingAssistant(context)
 }