Gemini API × Android Jetpack Compose Complete Integration Guide — Production Design Patterns for Kotlin Native AI Apps

// di/AiModule.kt
@Module
@InstallIn(SingletonComponent::class)
object AiModule {
 
    @Provides
    @Singleton
    fun provideGeminiDataSource(): GeminiDataSource = GeminiDataSource()
 
    @Provides
    @Singleton
    fun provideGeminiRepository(
        geminiDataSource: GeminiDataSource,
        chatDao: ChatDao
    ): GeminiRepository = GeminiRepositoryImpl(geminiDataSource, chatDao)
}
 
// di/DatabaseModule.kt
@Module
@InstallIn(SingletonComponent::class)
object DatabaseModule {
 
    @Provides
    @Singleton
    fun provideDatabase(@ApplicationContext context: Context): ChatDatabase {
        return Room.databaseBuilder(
            context,
            ChatDatabase::class.java,
            "chat_database"
        )
            .fallbackToDestructiveMigration()
            .build()
    }
 
    @Provides
    fun provideChatDao(database: ChatDatabase): ChatDao = database.chatDao()
}

// data/repository/GeminiRepository.kt
interface GeminiRepository {
    fun sendMessage(messages: List<ChatMessage>, imageBitmap: Bitmap? = null): Flow<GenerationResult>
    suspend fun getChatHistory(sessionId: String): List<ChatMessage>
    suspend fun saveChatMessage(message: ChatMessage)
    suspend fun clearSession(sessionId: String)
}
 
data class GenerationResult(
    val text: String = "",
    val isComplete: Boolean = false,
    val error: String? = null
)
 
// data/repository/GeminiRepositoryImpl.kt
class GeminiRepositoryImpl @Inject constructor(
    private val geminiDataSource: GeminiDataSource,
    private val chatDao: ChatDao
) : GeminiRepository {
 
    override fun sendMessage(
        messages: List<ChatMessage>,
        imageBitmap: Bitmap?
    ): Flow<GenerationResult> {
        val latestMessage = messages.lastOrNull()?.text ?: return emptyFlow()
        return geminiDataSource.streamChatResponse(messages, latestMessage, imageBitmap)
    }
 
    override suspend fun getChatHistory(sessionId: String): List<ChatMessage> {
        return chatDao.getSessionMessages(sessionId).map { it.toChatMessage() }
    }
 
    override suspend fun saveChatMessage(message: ChatMessage) {
        chatDao.insertMessage(message.toEntity())
    }
 
    override suspend fun clearSession(sessionId: String) {
        chatDao.clearSession(sessionId)
    }
}

// data/datasource/GeminiDataSource.kt
class GeminiDataSource @Inject constructor() {
 
    private fun buildGenerativeModel(modelName: String = "gemini-2.0-flash"): GenerativeModel {
        return GenerativeModel(
            modelName = modelName,
            apiKey = BuildConfig.GEMINI_API_KEY,
            generationConfig = generationConfig {
                temperature = 0.7f
                maxOutputTokens = 8192
                topP = 0.95f
            },
            safetySettings = listOf(
                SafetySetting(HarmCategory.HARASSMENT, BlockThreshold.MEDIUM_AND_ABOVE),
                SafetySetting(HarmCategory.HATE_SPEECH, BlockThreshold.MEDIUM_AND_ABOVE),
                SafetySetting(HarmCategory.DANGEROUS_CONTENT, BlockThreshold.MEDIUM_AND_ABOVE)
            )
        )
    }
 
    fun streamChatResponse(
        chatHistory: List<ChatMessage>,
        newMessage: String,
        imageBitmap: Bitmap? = null
    ): Flow<GenerationResult> = flow {
        val model = buildGenerativeModel()
 
        // Build history excluding the latest message
        val chat = model.startChat(
            history = chatHistory.dropLast(1).map { msg ->
                content(msg.role.name.lowercase()) { text(msg.text) }
            }
        )
 
        val content = content("user") {
            imageBitmap?.let { image(it) }
            text(newMessage)
        }
 
        try {
            chat.sendMessageStream(content).collect { chunk ->
                chunk.text?.let { deltaText ->
                    emit(GenerationResult(text = deltaText, isComplete = false))
                }
            }
            emit(GenerationResult(text = "", isComplete = true))
        } catch (e: GoogleGenerativeAIException) {
            emit(GenerationResult(error = "API error: ${e.message}", isComplete = true))
        } catch (e: Exception) {
            emit(GenerationResult(error = "Unexpected error: ${e.message}", isComplete = true))
        }
    }.flowOn(Dispatchers.IO)
}

// ui/chat/ChatViewModel.kt
@HiltViewModel
class ChatViewModel @Inject constructor(
    private val geminiRepository: GeminiRepository
) : ViewModel() {
 
    private val _uiState = MutableStateFlow(ChatUiState())
    val uiState: StateFlow<ChatUiState> = _uiState.asStateFlow()
 
    private var streamingJob: Job? = null
    val sessionId: String = UUID.randomUUID().toString()
 
    fun sendMessage(userInput: String, imageBitmap: Bitmap? = null) {
        if (userInput.isBlank() && imageBitmap == null) return
 
        val userMessage = ChatMessage(
            text = userInput,
            role = Role.USER,
            timestamp = System.currentTimeMillis()
        )
 
        _uiState.update { current ->
            current.copy(
                messages = current.messages + userMessage,
                isLoading = true,
                pendingBotText = "",
                errorMessage = null
            )
        }
 
        streamingJob = viewModelScope.launch {
            geminiRepository.sendMessage(
                messages = _uiState.value.messages,
                imageBitmap = imageBitmap
            ).collect { result ->
                when {
                    result.error != null -> handleError(result.error)
                    result.isComplete -> finalizeMessage()
                    else -> appendStreamingText(result.text)
                }
            }
        }
    }
 
    fun cancelStreaming() {
        streamingJob?.cancel()
        _uiState.update { it.copy(isLoading = false, pendingBotText = "") }
    }
 
    private fun appendStreamingText(text: String) {
        _uiState.update { current ->
            current.copy(pendingBotText = current.pendingBotText + text)
        }
    }
 
    private suspend fun finalizeMessage() {
        val finalText = _uiState.value.pendingBotText
        val botMessage = ChatMessage(
            text = finalText,
            role = Role.MODEL,
            timestamp = System.currentTimeMillis()
        )
        _uiState.update { current ->
            current.copy(
                messages = current.messages + botMessage,
                isLoading = false,
                pendingBotText = ""
            )
        }
        geminiRepository.saveChatMessage(botMessage)
    }
 
    private fun handleError(error: String) {
        _uiState.update { it.copy(isLoading = false, errorMessage = error) }
    }
}
 
data class ChatUiState(
    val messages: List<ChatMessage> = emptyList(),
    val isLoading: Boolean = false,
    val pendingBotText: String = "",
    val errorMessage: String? = null
)

// ui/chat/ChatScreen.kt
@Composable
fun ChatScreen(viewModel: ChatViewModel = hiltViewModel()) {
    val uiState by viewModel.uiState.collectAsStateWithLifecycle()
    val listState = rememberLazyListState()
    var inputText by remember { mutableStateOf("") }
    var selectedBitmap by remember { mutableStateOf<Bitmap?>(null) }
 
    // Image picker launcher
    val imagePickerLauncher = rememberLauncherForActivityResult(
        contract = ActivityResultContracts.GetContent()
    ) { uri ->
        uri?.let { /* Convert to Bitmap and store */ }
    }
 
    // Auto-scroll to latest content
    LaunchedEffect(uiState.messages.size, uiState.pendingBotText.length) {
        val totalItems = uiState.messages.size + if (uiState.pendingBotText.isNotEmpty()) 1 else 0
        if (totalItems > 0) listState.animateScrollToItem(totalItems - 1)
    }
 
    Scaffold(
        bottomBar = {
            ChatInputBar(
                inputText = inputText,
                selectedBitmap = selectedBitmap,
                isLoading = uiState.isLoading,
                onInputChange = { inputText = it },
                onImagePick = { imagePickerLauncher.launch("image/*") },
                onClearImage = { selectedBitmap = null },
                onSend = {
                    viewModel.sendMessage(inputText, selectedBitmap)
                    inputText = ""
                    selectedBitmap = null
                },
                onCancel = { viewModel.cancelStreaming() }
            )
        }
    ) { paddingValues ->
        LazyColumn(
            state = listState,
            modifier = Modifier.fillMaxSize().padding(paddingValues),
            contentPadding = PaddingValues(16.dp),
            verticalArrangement = Arrangement.spacedBy(8.dp)
        ) {
            items(uiState.messages, key = { it.timestamp }) { message ->
                ChatBubble(message = message)
            }
            if (uiState.pendingBotText.isNotEmpty()) {
                item(key = "streaming") { StreamingBubble(text = uiState.pendingBotText) }
            }
        }
    }
}
 
@Composable
fun StreamingBubble(text: String) {
    val cursorVisible by rememberInfiniteTransition(label = "cursor")
        .animateFloat(
            initialValue = 0f, targetValue = 1f,
            animationSpec = infiniteRepeatable(
                animation = tween(600),
                repeatMode = RepeatMode.Reverse
            ),
            label = "cursorAlpha"
        ).let { derivedStateOf { it.value > 0.5f } }
 
    Box(
        modifier = Modifier
            .fillMaxWidth(0.85f)
            .clip(RoundedCornerShape(12.dp))
            .background(MaterialTheme.colorScheme.secondaryContainer)
            .padding(12.dp)
    ) {
        Text(
            text = buildAnnotatedString {
                append(text)
                withStyle(SpanStyle(color = if (cursorVisible) LocalContentColor.current else Color.Transparent)) {
                    append("│")
                }
            },
            style = MaterialTheme.typography.bodyMedium
        )
    }
}

// ui/chat/ImagePickerHandler.kt
@Composable
fun rememberImagePickerHandler(
    onImageSelected: (Bitmap) -> Unit
): ImagePickerState {
    val context = LocalContext.current
 
    // Gallery picker
    val galleryLauncher = rememberLauncherForActivityResult(
        ActivityResultContracts.GetContent()
    ) { uri ->
        uri?.let {
            val bitmap = if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.P) {
                ImageDecoder.decodeBitmap(ImageDecoder.createSource(context.contentResolver, uri))
            } else {
                @Suppress("DEPRECATION")
                MediaStore.Images.Media.getBitmap(context.contentResolver, uri)
            }
            onImageSelected(bitmap)
        }
    }
 
    // Camera capture
    var tempUri by remember { mutableStateOf<Uri?>(null) }
    val cameraLauncher = rememberLauncherForActivityResult(
        ActivityResultContracts.TakePicture()
    ) { success ->
        if (success) {
            tempUri?.let { uri ->
                val bitmap = if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.P) {
                    ImageDecoder.decodeBitmap(ImageDecoder.createSource(context.contentResolver, uri))
                } else {
                    @Suppress("DEPRECATION")
                    MediaStore.Images.Media.getBitmap(context.contentResolver, uri)
                }
                onImageSelected(bitmap)
            }
        }
    }
 
    return ImagePickerState(
        onPickFromGallery = { galleryLauncher.launch("image/*") },
        onTakePhoto = {
            // Create temp file in cache dir for camera output
            val tempFile = File.createTempFile("camera_", ".jpg", context.cacheDir)
            val uri = FileProvider.getUriForFile(
                context,
                "${context.packageName}.fileprovider",
                tempFile
            )
            tempUri = uri
            cameraLauncher.launch(uri)
        }
    )
}
 
data class ImagePickerState(
    val onPickFromGallery: () -> Unit,
    val onTakePhoto: () -> Unit
)

// ui/chat/ChatInputBar.kt
@Composable
fun ChatInputBar(
    inputText: String,
    selectedBitmap: Bitmap?,
    isLoading: Boolean,
    onInputChange: (String) -> Unit,
    onSend: () -> Unit,
    onCancel: () -> Unit,
    onClearImage: () -> Unit,
    onPickImage: () -> Unit
) {
    Column(
        modifier = Modifier
            .fillMaxWidth()
            .background(MaterialTheme.colorScheme.surface)
            .padding(horizontal = 16.dp, vertical = 8.dp)
    ) {
        // Image preview thumbnail
        if (selectedBitmap != null) {
            Box(modifier = Modifier.padding(bottom = 8.dp)) {
                Image(
                    bitmap = selectedBitmap.asImageBitmap(),
                    contentDescription = "Selected image",
                    modifier = Modifier
                        .size(80.dp)
                        .clip(RoundedCornerShape(8.dp)),
                    contentScale = ContentScale.Crop
                )
                IconButton(
                    onClick = onClearImage,
                    modifier = Modifier
                        .size(24.dp)
                        .align(Alignment.TopEnd)
                        .background(MaterialTheme.colorScheme.error, CircleShape)
                ) {
                    Icon(
                        Icons.Default.Close,
                        contentDescription = "Remove image",
                        tint = Color.White,
                        modifier = Modifier.size(16.dp)
                    )
                }
            }
        }
 
        Row(
            modifier = Modifier.fillMaxWidth(),
            verticalAlignment = Alignment.CenterVertically,
            horizontalArrangement = Arrangement.spacedBy(8.dp)
        ) {
            // Image attachment button
            if (!isLoading) {
                IconButton(onClick = onPickImage) {
                    Icon(Icons.Default.AttachFile, contentDescription = "Attach image")
                }
            }
 
            // Text input
            OutlinedTextField(
                value = inputText,
                onValueChange = onInputChange,
                modifier = Modifier.weight(1f),
                placeholder = { Text("Type a message...") },
                maxLines = 4,
                keyboardOptions = KeyboardOptions(imeAction = ImeAction.Send),
                keyboardActions = KeyboardActions(onSend = { if (!isLoading) onSend() })
            )
 
            // Send / Cancel button
            if (isLoading) {
                IconButton(
                    onClick = onCancel,
                    modifier = Modifier.testTag("CancelButton")
                ) {
                    Icon(Icons.Default.Stop, contentDescription = "Cancel")
                }
            } else {
                IconButton(
                    onClick = onSend,
                    enabled = inputText.isNotBlank() || selectedBitmap != null,
                    modifier = Modifier.testTag("SendButton")
                ) {
                    Icon(Icons.AutoMirrored.Filled.Send, contentDescription = "Send")
                }
            }
        }
    }
}

<!-- AndroidManifest.xml -->
<provider
    android:name="androidx.core.content.FileProvider"
    android:authorities="${applicationId}.fileprovider"
    android:exported="false"
    android:grantUriPermissions="true">
    <meta-data
        android:name="android.support.FILE_PROVIDER_PATHS"
        android:resource="@xml/file_paths" />
</provider>

<!-- res/xml/file_paths.xml -->
<paths>
    <cache-path name="camera_images" path="." />
</paths>

// data/local/ChatEntity.kt
@Entity(tableName = "chat_messages")
data class ChatEntity(
    @PrimaryKey(autoGenerate = true) val id: Long = 0,
    val sessionId: String,
    val text: String,
    val role: String,       // "USER" or "MODEL"
    val imageUri: String?,
    val timestamp: Long,
    val tokenCount: Int = 0
)
 
// data/local/ChatDao.kt
@Dao
interface ChatDao {
    @Query("SELECT * FROM chat_messages WHERE sessionId = :sessionId ORDER BY timestamp ASC")
    fun observeSessionMessages(sessionId: String): Flow<List<ChatEntity>>
 
    @Query("SELECT * FROM chat_messages WHERE sessionId = :sessionId ORDER BY timestamp ASC")
    suspend fun getSessionMessages(sessionId: String): List<ChatEntity>
 
    @Insert(onConflict = OnConflictStrategy.REPLACE)
    suspend fun insertMessage(entity: ChatEntity)
 
    @Query("DELETE FROM chat_messages WHERE sessionId = :sessionId")
    suspend fun clearSession(sessionId: String)
 
    @Query("DELETE FROM chat_messages WHERE timestamp < :cutoffMs")
    suspend fun deleteOlderThan(cutoffMs: Long)
 
    @Query("SELECT COUNT(*) FROM chat_messages WHERE sessionId = :sessionId")
    suspend fun countMessages(sessionId: String): Int
}
 
// data/local/ChatDatabase.kt
@Database(entities = [ChatEntity::class], version = 1, exportSchema = false)
abstract class ChatDatabase : RoomDatabase() {
    abstract fun chatDao(): ChatDao
}

// worker/SummarizeWorker.kt
@HiltWorker
class SummarizeWorker @AssistedInject constructor(
    @Assisted context: Context,
    @Assisted params: WorkerParameters,
    private val geminiRepository: GeminiRepository
) : CoroutineWorker(context, params) {
 
    override suspend fun doWork(): Result {
        val sessionId = inputData.getString(KEY_SESSION_ID) ?: return Result.failure()
 
        return try {
            val messages = geminiRepository.getChatHistory(sessionId)
            if (messages.isEmpty()) return Result.success(workDataOf(KEY_SUMMARY to ""))
 
            val prompt = buildString {
                appendLine("Summarize the following conversation in 3 concise sentences:")
                messages.forEach { appendLine("${it.role}: ${it.text}") }
            }
 
            var summaryText = ""
            geminiRepository.sendMessage(
                listOf(ChatMessage(text = prompt, role = Role.USER))
            ).collect { result ->
                if (!result.isComplete && result.error == null) {
                    summaryText += result.text
                }
            }
 
            Result.success(workDataOf(KEY_SUMMARY to summaryText))
        } catch (e: Exception) {
            if (runAttemptCount < 3) Result.retry() else Result.failure()
        }
    }
 
    companion object {
        const val KEY_SESSION_ID = "session_id"
        const val KEY_SUMMARY = "summary"
 
        fun buildRequest(sessionId: String): OneTimeWorkRequest {
            return OneTimeWorkRequestBuilder<SummarizeWorker>()
                .setInputData(workDataOf(KEY_SESSION_ID to sessionId))
                .setConstraints(
                    Constraints.Builder()
                        .setRequiredNetworkType(NetworkType.CONNECTED)
                        .build()
                )
                .setBackoffCriteria(BackoffPolicy.EXPONENTIAL, 30, TimeUnit.SECONDS)
                .addTag("ai_summarize")
                .build()
        }
    }
}

// In ViewModel
fun requestSummary(sessionId: String) {
    val request = SummarizeWorker.buildRequest(sessionId)
    WorkManager.getInstance(context).enqueue(request)
 
    WorkManager.getInstance(context)
        .getWorkInfoByIdLiveData(request.id)
        .observeForever { info ->
            if (info?.state == WorkInfo.State.SUCCEEDED) {
                val summary = info.outputData.getString(SummarizeWorker.KEY_SUMMARY)
                // Update UI with summary
            }
        }
}

// In GeminiDataSource
suspend fun countTokens(messages: List<ChatMessage>): Int {
    val model = buildGenerativeModel()
    val contents = messages.map { msg ->
        content(msg.role.name.lowercase()) { text(msg.text) }
    }
    return try {
        model.countTokens(*contents.toTypedArray()).totalTokens
    } catch (e: Exception) {
        -1  // Return -1 to indicate unavailable, caller handles gracefully
    }
}

private suspend fun checkTokenCount() {
    val count = geminiDataSource.countTokens(_uiState.value.messages)
    if (count > 800_000) {  // 80% of 1M limit
        _uiState.update { it.copy(showTokenWarning = true) }
    }
}

fun Bitmap.downscaleForApi(maxDimension: Int = 1024): Bitmap {
    val ratio = maxDimension.toFloat() / maxOf(width, height)
    if (ratio >= 1f) return this
    val newWidth = (width * ratio).toInt()
    val newHeight = (height * ratio).toInt()
    return Bitmap.createScaledBitmap(this, newWidth, newHeight, true)
}

fun sendMessage(userInput: String, rawBitmap: Bitmap? = null) {
    val scaledBitmap = rawBitmap?.downscaleForApi(maxDimension = 1024)
    // ... rest of sendMessage logic
}

items(uiState.messages, key = { message -> message.timestamp }) { message ->
    ChatBubble(message = message)
}

class GeminiDataSource @Inject constructor() {
    private val apiKey: String
        get() = if (BuildConfig.DEBUG) {
            BuildConfig.GEMINI_API_KEY
        } else {
            "" // Production: use token from backend auth response
        }
}

// test/fake/FakeGeminiRepository.kt
class FakeGeminiRepository : GeminiRepository {
    var streamResponse: Flow<GenerationResult> = emptyFlow()
    val savedMessages = mutableListOf<ChatMessage>()
 
    override fun sendMessage(messages: List<ChatMessage>, imageBitmap: Bitmap?): Flow<GenerationResult> {
        return streamResponse
    }
 
    override suspend fun getChatHistory(sessionId: String): List<ChatMessage> = savedMessages
 
    override suspend fun saveChatMessage(message: ChatMessage) {
        savedMessages.add(message)
    }
 
    override suspend fun clearSession(sessionId: String) {
        savedMessages.clear()
    }
}
 
// test/ChatViewModelTest.kt
@OptIn(ExperimentalCoroutinesApi::class)
class ChatViewModelTest {
 
    @get:Rule
    val mainDispatcherRule = MainDispatcherRule()
 
    private val fakeRepository = FakeGeminiRepository()
    private lateinit var viewModel: ChatViewModel
 
    @Before
    fun setup() {
        viewModel = ChatViewModel(fakeRepository)
    }
 
    @Test
    fun `sendMessage appends user message immediately`() = runTest {
        fakeRepository.streamResponse = flow { emit(GenerationResult(text = "", isComplete = true)) }
        viewModel.sendMessage("Hello Gemini")
        val firstMessage = viewModel.uiState.value.messages.first()
        assertEquals("Hello Gemini", firstMessage.text)
        assertEquals(Role.USER, firstMessage.role)
    }
 
    @Test
    fun `streaming chunks accumulate in pendingBotText`() = runTest {
        fakeRepository.streamResponse = flowOf(
            GenerationResult(text = "Hello", isComplete = false),
            GenerationResult(text = " there", isComplete = false),
            GenerationResult(text = "", isComplete = true)
        )
        viewModel.sendMessage("Hi")
        advanceUntilIdle()
 
        val finalMessages = viewModel.uiState.value.messages
        assertEquals("Hello there", finalMessages.last().text)
        assertFalse(viewModel.uiState.value.isLoading)
    }
 
    @Test
    fun `error from repository surfaces in uiState`() = runTest {
        fakeRepository.streamResponse = flowOf(
            GenerationResult(error = "Rate limit exceeded", isComplete = true)
        )
        viewModel.sendMessage("test")
        advanceUntilIdle()
 
        assertNotNull(viewModel.uiState.value.errorMessage)
        assertFalse(viewModel.uiState.value.isLoading)
    }
}

// androidTest/ChatScreenTest.kt
@HiltAndroidTest
class ChatScreenTest {
 
    @get:Rule(order = 0)
    val hiltRule = HiltAndroidRule(this)
 
    @get:Rule(order = 1)
    val composeTestRule = createAndroidComposeRule<MainActivity>()
 
    @Test
    fun chatInput_sendMessage_displaysUserBubble() {
        composeTestRule.onNodeWithTag("ChatInput")
            .performTextInput("What can Gemini do?")
        composeTestRule.onNodeWithTag("SendButton")
            .performClick()
        composeTestRule.onNodeWithText("What can Gemini do?")
            .assertIsDisplayed()
    }
 
    @Test
    fun cancelButton_appearsWhileLoading() {
        composeTestRule.onNodeWithTag("ChatInput")
            .performTextInput("Generate a long response")
        composeTestRule.onNodeWithTag("SendButton").performClick()
 
        // Cancel button should appear during streaming
        composeTestRule.onNodeWithTag("CancelButton").assertIsDisplayed()
    }
}

// GeminiFailure.kt — normalize failures into operable categories
enum class GeminiFailureType {
    RATE_LIMIT, SAFETY_BLOCKED, RECITATION, TIMEOUT, AUTH, NETWORK, UNKNOWN
}
 
fun classifyGeminiFailure(t: Throwable): GeminiFailureType = when {
    t is com.google.ai.client.generativeai.type.PromptBlockedException -> GeminiFailureType.SAFETY_BLOCKED
    t is com.google.ai.client.generativeai.type.ServerException &&
        t.message?.contains("429") == true -> GeminiFailureType.RATE_LIMIT
    t.message?.contains("RECITATION") == true -> GeminiFailureType.RECITATION
    t is java.net.SocketTimeoutException -> GeminiFailureType.TIMEOUT
    t is com.google.ai.client.generativeai.type.InvalidAPIKeyException -> GeminiFailureType.AUTH
    t is java.io.IOException -> GeminiFailureType.NETWORK
    else -> GeminiFailureType.UNKNOWN
}

// Called from the catch block in GeminiDataSource.kt
import com.google.firebase.crashlytics.FirebaseCrashlytics
 
private fun reportFailure(t: Throwable, modelName: String, attempt: Int) {
    val type = classifyGeminiFailure(t)
    FirebaseCrashlytics.getInstance().apply {
        setCustomKey("gemini_failure_type", type.name)
        setCustomKey("gemini_model", modelName)
        setCustomKey("gemini_attempt", attempt)
        // Treat expected failures (safety / 429) as logs, not crashes,
        // so they don't drown the real NPEs you need to fix.
        if (type == GeminiFailureType.SAFETY_BLOCKED || type == GeminiFailureType.RATE_LIMIT) {
            log("gemini_${type.name.lowercase()} model=$modelName attempt=$attempt")
        } else {
            recordException(t)
        }
    }
}

class ChatViewModel @Inject constructor(
    private val repository: GeminiRepository,
    private val savedState: SavedStateHandle
) : ViewModel() {
 
    fun sendMessage(prompt: String) {
        viewModelScope.launch {
            val buffer = StringBuilder()
            repository.streamResponse(prompt)
                .onCompletion { cause ->
                    if (cause == null) {
                        repository.persistTurn(prompt, buffer.toString()) // commit to Room
                        savedState[KEY_PARTIAL] = ""                       // clear the stash
                    }
                }
                .collect { chunk ->
                    buffer.append(chunk.text)
                    savedState[KEY_PARTIAL] = buffer.toString()            // stash each chunk
                    emitToUi(buffer.toString())
                }
        }
    }
 
    init {
        // On return, restore any stashed text as an "incomplete" response
        val pending = savedState.get<String>(KEY_PARTIAL).orEmpty()
        if (pending.isNotBlank()) {
            emitToUi(pending, incomplete = true)  // "interrupted — regenerate?" UI
        }
    }
 
    companion object { private const val KEY_PARTIAL = "partial_response" }
}