Improve Mistral rate-limit handling and add cooldown probe scripts

app/src/main/kotlin/com/google/ai/sample/ScreenCaptureApiClients.kt

@@ -15,6 +15,7 @@ import kotlinx.serialization.json.JsonClassDiscriminator
 import kotlinx.serialization.modules.SerializersModule
 import kotlinx.serialization.modules.polymorphic
 import kotlinx.serialization.modules.subclass
+import com.google.ai.sample.network.MistralRequestCoordinator
 import okhttp3.MediaType.Companion.toMediaType
 import okhttp3.OkHttpClient
 import okhttp3.Request
@@ -70,7 +71,7 @@ data class ServiceMistralResponseMessage(
     val content: String
 )
 
-internal suspend fun callMistralApi(modelName: String, apiKey: String, chatHistory: List<Content>, inputContent: Content): Pair<String?, String?> {
+internal suspend fun callMistralApi(modelName: String, apiKeys: List<String>, chatHistory: List<Content>, inputContent: Content): Pair<String?, String?> {
     var responseText: String? = null
     var errorMessage: String? = null
 
@@ -126,10 +127,18 @@ internal suspend fun callMistralApi(modelName: String, apiKey: String, chatHisto
             .url("https://api.mistral.ai/v1/chat/completions")
             .post(jsonBody.toRequestBody(mediaType))
             .addHeader("Content-Type", "application/json")
-            .addHeader("Authorization", "Bearer $apiKey")
+            .addHeader("Authorization", "Bearer ${apiKeys.first()}")
             .build()
 
-        client.newCall(request).execute().use { response ->
+        val coordinated = MistralRequestCoordinator.execute(apiKeys = apiKeys, maxAttempts = apiKeys.size * 4 + 8) { key ->
+            client.newCall(
+                request.newBuilder()
+                    .header("Authorization", "Bearer $key")
+                    .build()
+            ).execute()
+        }
+
+        coordinated.response.use { response ->
             val responseBody = response.body?.string()
             if (!response.isSuccessful) {
                 Log.e("ScreenCaptureService", "Mistral API Error ($response.code): $responseBody")

app/src/main/kotlin/com/google/ai/sample/ScreenCaptureService.kt

@@ -297,7 +297,11 @@ class ScreenCaptureService : Service() {
                             if (apiProvider == ApiProvider.VERCEL) {
                                 responseText = callVercelApi(applicationContext, modelName, apiKey, chatHistoryDtos, inputContentDto)
                             } else if (apiProvider == ApiProvider.MISTRAL) {
-                                val result = callMistralApi(modelName, apiKey, chatHistory, inputContent)
+                                val apiKeyManager = ApiKeyManager.getInstance(applicationContext)
+                                val availableKeys = apiKeyManager.getApiKeys(ApiProvider.MISTRAL)
+                                    .filter { it.isNotBlank() }
+                                    .distinct()
+                                val result = callMistralApi(modelName, availableKeys, chatHistory, inputContent)
                                 responseText = result.first
                                 errorMessage = result.second
                             } else if (apiProvider == ApiProvider.PUTER) {

app/src/main/kotlin/com/google/ai/sample/feature/multimodal/PhotoReasoningViewModel.kt

@@ -34,6 +34,7 @@ import com.google.ai.sample.feature.multimodal.ModelDownloadManager
 import com.google.ai.sample.ModelOption
 import com.google.ai.sample.GenerativeAiViewModelFactory
 import com.google.ai.sample.InferenceBackend
+import com.google.ai.sample.network.MistralRequestCoordinator
 import com.google.ai.sample.feature.multimodal.dtos.toDto
 import com.google.ai.sample.feature.multimodal.dtos.TempFilePathCollector
 import kotlinx.coroutines.Dispatchers
@@ -70,7 +71,6 @@ import kotlinx.serialization.modules.subclass
 import com.google.ai.sample.webrtc.WebRTCSender
 import com.google.ai.sample.webrtc.SignalingClient
 import org.webrtc.IceCandidate
-import kotlin.math.max
 
 class PhotoReasoningViewModel(
     application: Application,
@@ -183,11 +183,14 @@ class PhotoReasoningViewModel(
     // to avoid re-executing already-executed commands
     private var incrementalCommandCount = 0
 
-    // Mistral rate limiting per API key (1.1 seconds between requests with same key)
-    private val mistralNextAllowedRequestAtMsByKey = mutableMapOf<String, Long>()
-    private var lastMistralTokenTimeMs = 0L
-    private var lastMistralTokenKey: String? = null
-    private val MISTRAL_MIN_INTERVAL_MS = 1100L
+    private data class QueuedMistralScreenshotRequest(
+        val bitmap: Bitmap,
+        val screenshotUri: String,
+        val screenInfo: String?
+    )
+    private val mistralAutoScreenshotQueueLock = Any()
+    private var mistralAutoScreenshotInFlight = false
+    private var queuedMistralScreenshotRequest: QueuedMistralScreenshotRequest? = null
 
     // Accumulated full text during streaming for incremental command parsing
     private var streamingAccumulatedText = StringBuilder()
@@ -609,6 +612,7 @@ class PhotoReasoningViewModel(
         val currentModel = com.google.ai.sample.GenerativeAiViewModelFactory.getCurrentModel()
 
         clearStaleErrorState()
+        stopExecutionFlag.set(false)
 
         // Check for Human Expert model
         if (currentModel == ModelOption.HUMAN_EXPERT) {
@@ -1024,15 +1028,16 @@ class PhotoReasoningViewModel(
         )
     }
 
-private fun reasonWithMistral(
-    userInput: String,
-    selectedImages: List<Bitmap>,
-    screenInfoForPrompt: String? = null,
-    imageUrisForChat: List<String>? = null
-) {
-    _uiState.value = PhotoReasoningUiState.Loading
-    val context = appContext
-    val apiKeyManager = ApiKeyManager.getInstance(context)
+    private fun reasonWithMistral(
+        userInput: String,
+        selectedImages: List<Bitmap>,
+        screenInfoForPrompt: String? = null,
+        imageUrisForChat: List<String>? = null
+    ) {
+        _uiState.value = PhotoReasoningUiState.Loading
+        _showStopNotificationFlow.value = true
+        val context = appContext
+        val apiKeyManager = ApiKeyManager.getInstance(context)
 
     val initialApiKey = apiKeyManager.getCurrentApiKey(ApiProvider.MISTRAL)
     if (initialApiKey.isNullOrEmpty()) {
@@ -1054,7 +1059,8 @@ private fun reasonWithMistral(
 
     resetStreamingCommandState()
 
-    viewModelScope.launch(Dispatchers.IO) {
+    currentReasoningJob?.cancel()
+    currentReasoningJob = viewModelScope.launch(Dispatchers.IO) {
         try {
             val currentModel = com.google.ai.sample.GenerativeAiViewModelFactory.getCurrentModel()
             val genSettings = com.google.ai.sample.util.GenerationSettingsPreferences.loadSettings(context, currentModel.modelName)
@@ -1132,124 +1138,32 @@ private fun reasonWithMistral(
 
             // Validate that we have at least one key before proceeding
             require(availableKeys.isNotEmpty()) { "No valid Mistral API keys available after filtering" }
-
-            fun markKeyCooldown(key: String, referenceTimeMs: Long) {
-                val nextAllowedAt = referenceTimeMs + MISTRAL_MIN_INTERVAL_MS
-                val existing = mistralNextAllowedRequestAtMsByKey[key] ?: 0L
-                mistralNextAllowedRequestAtMsByKey[key] = max(existing, nextAllowedAt)
-            }
-
-            fun remainingWaitForKeyMs(key: String, nowMs: Long): Long {
-                val nextAllowedAt = mistralNextAllowedRequestAtMsByKey[key] ?: 0L
-                return (nextAllowedAt - nowMs).coerceAtLeast(0L)
-            }
-
-            fun isRetryableMistralFailure(code: Int): Boolean {
-                return code == 429 || code >= 500
-            }
-
-            var response: okhttp3.Response? = null
-            var selectedKeyForResponse: String? = null
-            var consecutiveFailures = 0
-            var blockedKeysThisRound = mutableSetOf<String>()
-
-            val maxAttempts = availableKeys.size * 2 + 3 // Allow cycling through all keys at least twice
-            while (response == null && consecutiveFailures < maxAttempts) {
-                if (stopExecutionFlag.get()) break
-
-                val now = System.currentTimeMillis()
-                val keyPool = availableKeys.filter { it !in blockedKeysThisRound }.ifEmpty {
-                    blockedKeysThisRound.clear()
-                    availableKeys
-                }
-
-                val keyWithLeastWait = keyPool.minByOrNull { remainingWaitForKeyMs(it, now) } ?: availableKeys.first()
-                val waitMs = remainingWaitForKeyMs(keyWithLeastWait, now)
-                if (waitMs > 0L) {
-                    delay(waitMs)
+            val maxAttempts = availableKeys.size * 4 + 8
+            val coordinated = MistralRequestCoordinator.execute(
+                apiKeys = availableKeys,
+                maxAttempts = maxAttempts
+            ) { selectedKey ->
+                if (stopExecutionFlag.get()) {
+                    throw IOException("Mistral request aborted.")
                 }
-
-                val selectedKey = keyWithLeastWait
-                selectedKeyForResponse = selectedKey
-
-                try {
-                    val attemptResponse = client.newCall(buildRequest(selectedKey)).execute()
-                    val requestEndMs = System.currentTimeMillis()
-                    markKeyCooldown(selectedKey, requestEndMs)
-
-                    if (attemptResponse.isSuccessful) {
-                        response = attemptResponse
-                        break
-                    }
-
-                    val isRetryable = isRetryableMistralFailure(attemptResponse.code)
-                    if (!isRetryable) {
-                        val errBody = attemptResponse.body?.string()
-                        attemptResponse.close()
-                        throw IllegalStateException("Mistral Error ${attemptResponse.code}: $errBody")
-                    }
-
-                    attemptResponse.close()
-                    blockedKeysThisRound.add(selectedKey)
-                    consecutiveFailures++
-                    withContext(Dispatchers.Main) {
-                        replaceAiMessageText(
-                            "Mistral temporär nicht verfügbar (Versuch $consecutiveFailures/$maxAttempts). Wiederhole...",
-                            isPending = true
-                        )
-                    }
-                } catch (e: IOException) {
-                    val requestEndMs = System.currentTimeMillis()
-                    markKeyCooldown(selectedKey, requestEndMs)
-                    blockedKeysThisRound.add(selectedKey)
-                    consecutiveFailures++
-                    if (consecutiveFailures >= 5) {
-                        throw IOException("Mistral request failed after 5 attempts: ${e.message}", e)
-                    }
-                    withContext(Dispatchers.Main) {
-                        replaceAiMessageText(
-                        if (consecutiveFailures >= maxAttempts) {
-                            throw IOException("Mistral request failed after $maxAttempts attempts: ${e.message}", e)
-                        )
-                    }
-                }
-                                "Mistral Netzwerkfehler (Versuch $consecutiveFailures/$maxAttempts). Wiederhole...",
-
-            if (stopExecutionFlag.get()) {
-                throw IOException("Mistral request aborted.")
+                client.newCall(buildRequest(selectedKey)).execute()
             }
-
-            val finalResponse = response ?: throw IOException("Mistral request failed after 5 attempts.")
+            val finalResponse = coordinated.response
 
             if (!finalResponse.isSuccessful) {
                 val errBody = finalResponse.body?.string()
                 finalResponse.close()
-            val finalResponse = response ?: throw IOException("Mistral request failed after $maxAttempts attempts.")
+                throw IOException("Mistral Error ${finalResponse.code}: $errBody")
             }
 
             val body = finalResponse.body ?: throw IOException("Empty response body from Mistral")
             val aiResponseText = openAiStreamParser.parse(body) { accText ->
-                selectedKeyForResponse?.let { key ->
-                    lastMistralTokenKey = key
-                    lastMistralTokenTimeMs = System.currentTimeMillis()
-                    markKeyCooldown(key, lastMistralTokenTimeMs)
-                } ?: run {
-                    Log.w(TAG, "selectedKeyForResponse is null during streaming callback")
-                }
                 withContext(Dispatchers.Main) {
                     replaceAiMessageText(accText, isPending = true)
                     processCommandsIncrementally(accText)
                 }
             }
             finalResponse.close()
-            selectedKeyForResponse?.let { key ->
-                val reference = if (lastMistralTokenKey == key && lastMistralTokenTimeMs > 0L) {
-                    lastMistralTokenTimeMs
-                } else {
-                    System.currentTimeMillis()
-                }
-                markKeyCooldown(key, reference)
-            }
 
             withContext(Dispatchers.Main) {
                 _uiState.value = PhotoReasoningUiState.Success(aiResponseText)
@@ -1261,9 +1175,15 @@ private fun reasonWithMistral(
             withContext(Dispatchers.Main) {
                 Log.e(TAG, "Mistral API call failed", e)
                 _uiState.value = PhotoReasoningUiState.Error(e.message ?: "Unknown error")
+                _chatState.replaceLastPendingMessage()
                 appendErrorMessage("Error: ${e.message}")
                 saveChatHistory(context)
             }
+        } finally {
+            withContext(Dispatchers.Main) {
+                releaseAndDrainMistralAutoScreenshotQueue()
+                refreshStopButtonState()
+            }
         }
     }
 }
@@ -2360,16 +2280,22 @@ private fun processCommands(text: String) {
                             _commandExecutionStatus.value = status
                         }
 
-                        // Create prompt with screen information if available
-                        val genericAnalysisPrompt = createGenericScreenshotPrompt()
-
-                        // Re-send the query with only the latest screenshot
-                        reason(
-                            userInput = genericAnalysisPrompt,
-                            selectedImages = listOf(bitmap),
-                            screenInfoForPrompt = screenInfo,
-                            imageUrisForChat = listOf(screenshotUri.toString()) // Add this argument
-                        )
+                        val currentModel = GenerativeAiViewModelFactory.getCurrentModel()
+                        if (currentModel.apiProvider == ApiProvider.MISTRAL) {
+                            enqueueMistralAutoScreenshotRequest(
+                                bitmap = bitmap,
+                                screenshotUri = screenshotUri.toString(),
+                                screenInfo = screenInfo
+                            )
+                        } else {
+                            // Re-send the query with only the latest screenshot
+                            reason(
+                                userInput = createGenericScreenshotPrompt(),
+                                selectedImages = listOf(bitmap),
+                                screenInfoForPrompt = screenInfo,
+                                imageUrisForChat = listOf(screenshotUri.toString())
+                            )
+                        }
 
                         PhotoReasoningScreenshotUiNotifier.showAddedToConversation(context)
                     } else {
@@ -2392,5 +2318,57 @@ private fun processCommands(text: String) {
             }
         }
     }
+
+    private fun enqueueMistralAutoScreenshotRequest(
+        bitmap: Bitmap,
+        screenshotUri: String,
+        screenInfo: String?
+    ) {
+        val request = QueuedMistralScreenshotRequest(
+            bitmap = bitmap,
+            screenshotUri = screenshotUri,
+            screenInfo = screenInfo
+        )
+        var shouldStartNow = false
+        synchronized(mistralAutoScreenshotQueueLock) {
+            if (mistralAutoScreenshotInFlight) {
+                queuedMistralScreenshotRequest = request
+                Log.d(TAG, "Mistral auto screenshot request queued (latest wins).")
+            } else {
+                mistralAutoScreenshotInFlight = true
+                shouldStartNow = true
+            }
+        }
+        if (shouldStartNow) {
+            dispatchMistralAutoScreenshotRequest(request)
+        }
+    }
+
+    private fun dispatchMistralAutoScreenshotRequest(request: QueuedMistralScreenshotRequest) {
+        val genericAnalysisPrompt = createGenericScreenshotPrompt()
+        reasonWithMistral(
+            userInput = genericAnalysisPrompt,
+            selectedImages = listOf(request.bitmap),
+            screenInfoForPrompt = request.screenInfo,
+            imageUrisForChat = listOf(request.screenshotUri)
+        )
+    }
+
+    private fun releaseAndDrainMistralAutoScreenshotQueue() {
+        val nextRequest: QueuedMistralScreenshotRequest? = synchronized(mistralAutoScreenshotQueueLock) {
+            val queued = queuedMistralScreenshotRequest
+            if (queued == null) {
+                mistralAutoScreenshotInFlight = false
+                null
+            } else {
+                queuedMistralScreenshotRequest = null
+                queued
+            }
+        }
+        if (nextRequest != null) {
+            Log.d(TAG, "Draining queued Mistral auto screenshot request.")
+            dispatchMistralAutoScreenshotRequest(nextRequest)
+        }
+    }
 
 }