improve the profile cache

composeApp/src/androidMain/kotlin/su/reya/coop/App.kt

@@ -6,11 +6,14 @@ import android.os.Build
 import androidx.activity.ComponentActivity
 import androidx.activity.compose.BackHandler
 import androidx.compose.foundation.isSystemInDarkTheme
+import androidx.compose.foundation.layout.Box
+import androidx.compose.foundation.layout.fillMaxSize
 import androidx.compose.material3.ExperimentalMaterial3Api
 import androidx.compose.material3.ExperimentalMaterial3ExpressiveApi
 import androidx.compose.material3.MaterialExpressiveTheme
 import androidx.compose.material3.MotionScheme
 import androidx.compose.material3.SnackbarHostState
+import androidx.compose.material3.Text
 import androidx.compose.material3.Typography
 import androidx.compose.material3.darkColorScheme
 import androidx.compose.material3.dynamicDarkColorScheme
@@ -24,6 +27,8 @@ import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.remember
 import androidx.compose.runtime.setValue
 import androidx.compose.runtime.staticCompositionLocalOf
+import androidx.compose.ui.Alignment
+import androidx.compose.ui.Modifier
 import androidx.compose.ui.platform.LocalContext
 import androidx.core.util.Consumer
 import androidx.lifecycle.ViewModel
@@ -230,26 +235,34 @@ fun App(
                         NewIdentityScreen(accountViewModel)
                     }
                     entry<Screen.Chat> { key ->
-                        val factory = remember(key) {
-                            object : ViewModelProvider.Factory {
-                                override fun <T : ViewModel> create(modelClass: Class<T>): T {
-                                    @Suppress("UNCHECKED_CAST")
-                                    return ChatScreenViewModel(
-                                        key.id,
-                                        key.screening,
-                                        accountRepository,
-                                        chatRepository
-                                    ) as T
+                        val initialRoom = remember(key.id) { chatRepository.getChatRoom(key.id) }
+                        
+                        if (initialRoom != null) {
+                            val factory = remember(initialRoom) {
+                                object : ViewModelProvider.Factory {
+                                    override fun <T : ViewModel> create(modelClass: Class<T>): T {
+                                        return ChatScreenViewModel(
+                                            initialRoom,
+                                            key.screening,
+                                            accountRepository,
+                                            chatRepository
+                                        ) as T
+                                    }
                                 }
                             }
+                            ChatScreen(
+                                viewModel<ChatScreenViewModel>(
+                                    key = key.id.toString(),
+                                    factory = factory
+                                ),
+                                accountViewModel
+                            )
+                        } else {
+                            // Handle the rare case where the room isn't in DB (e.g., invalid deep link)
+                            Box(Modifier.fillMaxSize(), contentAlignment = Alignment.Center) {
+                                Text("Room not found")
+                            }
                         }
-                        ChatScreen(
-                            viewModel<ChatScreenViewModel>(
-                                key = key.id.toString(),
-                                factory = factory
-                            ),
-                            accountViewModel
-                        )
                     }
                     entry<Screen.NewChat> {
                         NewChatScreen(accountViewModel, chatViewModel)

composeApp/src/androidMain/kotlin/su/reya/coop/screens/HomeScreen.kt

@@ -98,9 +98,9 @@ import su.reya.coop.RoomKind
 import su.reya.coop.RoomUiState
 import su.reya.coop.Screen
 import su.reya.coop.ago
+import su.reya.coop.flow
 import su.reya.coop.shared.Avatar
 import su.reya.coop.shared.getExpressiveFontFamily
-import su.reya.coop.uiStateFlow
 import su.reya.coop.viewmodel.AccountViewModel
 import su.reya.coop.viewmodel.ChatViewModel
 
@@ -620,9 +620,9 @@ fun NewRequests(requests: List<Room>) {
     val firstRoom = requests.getOrNull(0)
     val secondRoom = requests.getOrNull(1)
 
-    val firstRoomState by (firstRoom as Room).uiStateFlow(profileCache)
+    val firstRoomState by (firstRoom as Room).flow(profileCache)
         .collectAsStateWithLifecycle(RoomUiState())
-    val secondRoomState by (secondRoom ?: firstRoom).uiStateFlow(profileCache)
+    val secondRoomState by (secondRoom ?: firstRoom).flow(profileCache)
         .collectAsStateWithLifecycle(RoomUiState())
 
     val supportingText = when {
@@ -695,8 +695,7 @@ fun NewRequests(requests: List<Room>) {
 @Composable
 fun ChatRoom(room: Room, onClick: () -> Unit) {
     val profileCache = LocalProfileCache.current
-    val roomState by room.uiStateFlow(profileCache)
-        .collectAsStateWithLifecycle(RoomUiState())
+    val roomState by room.flow(profileCache).collectAsStateWithLifecycle(RoomUiState())
 
     ListItem(
         modifier = Modifier.clickable(onClick = onClick),

composeApp/src/androidMain/kotlin/su/reya/coop/screens/chat/ChatScreen.kt

@@ -94,12 +94,11 @@ import rust.nostr.sdk.UnsignedEvent
 import su.reya.coop.LocalNavigator
 import su.reya.coop.LocalProfileCache
 import su.reya.coop.LocalSnackbarHostState
-import su.reya.coop.Room
 import su.reya.coop.RoomUiState
 import su.reya.coop.Screen
+import su.reya.coop.flow
 import su.reya.coop.formatAsGroup
 import su.reya.coop.shared.Avatar
-import su.reya.coop.uiStateFlow
 import su.reya.coop.viewmodel.AccountViewModel
 import su.reya.coop.viewmodel.ChatScreenViewModel
 
@@ -118,25 +117,11 @@ fun ChatScreen(
     val scope = rememberCoroutineScope()
     val listState = rememberLazyListState()
 
-    val id = viewModel.id
     val currentUser by viewModel.currentUser.collectAsStateWithLifecycle()
-    val chatRooms by viewModel.chatRooms.collectAsStateWithLifecycle()
-    val room by remember(id) { derivedStateOf { chatRooms.firstOrNull { it.id == id } } }
+    val pubkey = currentUser?.publicKey
 
-    // Show empty screen
-    if (room == null) {
-        Box(
-            modifier = Modifier.fillMaxSize(),
-            contentAlignment = Alignment.Center
-        ) {
-            Text(
-                text = "Something went wrong.",
-                style = MaterialTheme.typography.titleMediumEmphasized,
-                color = MaterialTheme.colorScheme.onSurface
-            )
-        }
-        return
-    }
+    val room by viewModel.room.collectAsStateWithLifecycle()
+    val roomState by room.flow(profileCache, pubkey).collectAsStateWithLifecycle(RoomUiState())
 
     val loading = viewModel.loading
     val newOtherMessages = viewModel.newOtherMessages
@@ -146,9 +131,6 @@ fun ChatScreen(
     val groupedMessages =
         remember { derivedStateOf { messages.groupBy { it.createdAt().formatAsGroup() } } }
 
-    val roomState by (room as Room).uiStateFlow(profileCache, currentUser?.publicKey)
-        .collectAsStateWithLifecycle(RoomUiState())
-
     var text by remember { mutableStateOf("") }
     var selectedMessage by remember { mutableStateOf<Pair<MessageModel, Rect>?>(null) }
     var replyingTo by remember { mutableStateOf<MessageModel?>(null) }
@@ -193,9 +175,7 @@ fun ChatScreen(
         }
     }
 
-    Box(
-        modifier = Modifier.fillMaxSize()
-    ) {
+    Box(modifier = Modifier.fillMaxSize()) {
         Scaffold(
             modifier = Modifier.blur(blurAmount),
             contentWindowInsets = ScaffoldDefaults.contentWindowInsets.union(WindowInsets.ime),
@@ -207,7 +187,7 @@ fun ChatScreen(
                         Row(
                             verticalAlignment = Alignment.CenterVertically,
                             modifier = Modifier.clickable {
-                                room?.members?.firstOrNull()?.let { pubkey ->
+                                room.members.firstOrNull()?.let { pubkey ->
                                     navigator.navigate(Screen.Profile(pubkey.toBech32()))
                                 }
                             }
@@ -265,7 +245,7 @@ fun ChatScreen(
                             .padding(bottom = innerPadding.calculateBottomPadding())
                     ) {
                         if (requireScreening) {
-                            room?.let { ScreenerCard(accountViewModel, it) }
+                            ScreenerCard(accountViewModel, room)
                         }
 
                         when (messages.isNotEmpty()) {
@@ -283,8 +263,7 @@ fun ChatScreen(
                                             items = messagesInGroup,
                                             key = { it.ensureId().id()?.toHex()!! }
                                         ) { event ->
-                                            val model =
-                                                rememberMessageModel(event, currentUser?.publicKey)
+                                            val model = rememberMessageModel(event, pubkey)
 
                                             val replyPreview =
                                                 remember(model.replyEventIds, messages.size) {
@@ -298,7 +277,9 @@ fun ChatScreen(
                                                 modifier = Modifier.fillMaxWidth(),
                                                 verticalArrangement = Arrangement.spacedBy(2.dp)
                                             ) {
-                                                replyPreview?.let { ReplyPreview(it, model.isMine) }
+                                                replyPreview?.let {
+                                                    ReplyPreview(it, model.isMine)
+                                                }
                                                 ChatMessage(
                                                     model = model,
                                                     modifier = Modifier.graphicsLayer {

shared/src/commonMain/kotlin/su/reya/coop/Room.kt

@@ -71,7 +71,7 @@ data class RoomUiState(
     val isGroup: Boolean = false
 )
 
-fun Room.uiStateFlow(
+fun Room.flow(
     profileCache: ProfileCache,
     currentUser: PublicKey? = null
 ): Flow<RoomUiState> {

shared/src/commonMain/kotlin/su/reya/coop/viewmodel/AccountViewModel.kt

@@ -13,7 +13,6 @@ import su.reya.coop.repository.AccountState
 class AccountViewModel(
     private val repository: AccountRepository,
 ) : ViewModel(), ErrorHost by repository {
-
     val state: StateFlow<AccountState> = repository.state
     val isUpdatingProfile: StateFlow<Boolean> = repository.isUpdatingProfile
     val currentUserProfile: StateFlow<Profile?> = repository.currentUserProfile

shared/src/commonMain/kotlin/su/reya/coop/viewmodel/ChatScreenViewModel.kt

@@ -7,30 +7,40 @@ import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.setValue
 import androidx.lifecycle.ViewModel
 import androidx.lifecycle.viewModelScope
+import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.stateIn
 import kotlinx.coroutines.launch
 import rust.nostr.sdk.EventId
 import rust.nostr.sdk.UnsignedEvent
-import su.reya.coop.Profile
 import su.reya.coop.Room
 import su.reya.coop.repository.AccountRepository
 import su.reya.coop.repository.ChatRepository
 import su.reya.coop.roomId
 
 class ChatScreenViewModel(
-    val id: Long,
+    initialRoom: Room,
     screening: Boolean,
     accountRepository: AccountRepository,
     private val chatRepository: ChatRepository,
 ) : ViewModel(), ErrorHost by chatRepository {
-    val currentUser: StateFlow<Profile?> = accountRepository.currentUserProfile
-    val chatRooms: StateFlow<List<Room>> = chatRepository.chatRooms
-
     var loading by mutableStateOf(true)
     var newOtherMessages by mutableIntStateOf(0)
     var requireScreening by mutableStateOf(screening)
     val messages = mutableStateListOf<UnsignedEvent>()
 
+    val currentUser = accountRepository.currentUserProfile
+    val id = initialRoom.id
+    
+    val room: StateFlow<Room> = chatRepository.chatRooms
+        .map { rooms -> rooms.find { it.id == id } ?: initialRoom }
+        .stateIn(
+            scope = viewModelScope,
+            started = SharingStarted.WhileSubscribed(5000),
+            initialValue = initialRoom
+        )
+
     init {
         loadMessages()
         connect()

shared/src/commonMain/kotlin/su/reya/coop/viewmodel/ProfileCache.kt

@@ -1,16 +1,15 @@
 package su.reya.coop.viewmodel
 
+import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.SupervisorJob
 import kotlinx.coroutines.channels.Channel
-import kotlinx.coroutines.coroutineScope
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.asStateFlow
+import kotlinx.coroutines.flow.update
 import kotlinx.coroutines.launch
-import kotlinx.coroutines.sync.Mutex
-import kotlinx.coroutines.sync.withLock
 import kotlinx.coroutines.withTimeoutOrNull
 import rust.nostr.sdk.PublicKey
 import su.reya.coop.Profile
@@ -19,27 +18,41 @@ import kotlin.time.Duration.Companion.milliseconds
 
 class ProfileCache(
     private val nostr: Nostr,
+    private val defaultDispatcher: CoroutineDispatcher = Dispatchers.Default,
 ) : ErrorHost by createErrorHost() {
-    private val scope = CoroutineScope(SupervisorJob() + Dispatchers.Default)
-    private val profilesMutex = Mutex()
-    private val profiles = mutableMapOf<PublicKey, MutableStateFlow<Profile?>>()
+    private val scope = CoroutineScope(SupervisorJob() + defaultDispatcher)
+    private val profiles = MutableStateFlow<Map<PublicKey, MutableStateFlow<Profile?>>>(emptyMap())
+    private val seenPublicKeys = MutableStateFlow<Set<PublicKey>>(emptySet())
     private val metadataRequestChannel = Channel<PublicKey>(Channel.UNLIMITED)
-    private val seenPublicKeys = mutableSetOf<PublicKey>()
 
     init {
-        scope.launch { runObserver() }
-        scope.launch { runMetadataBatching() }
         scope.launch {
-            nostr.waitUntilInitialized()
-            loadCacheMetadata()
+            try {
+                runObserver()
+            } catch (e: Exception) {
+                showError("Metadata observer failed: ${e.message}")
+            }
+        }
+        scope.launch {
+            try {
+                runMetadataBatching()
+            } catch (e: Exception) {
+                showError("Metadata batching failed: ${e.message}")
+            }
+        }
+        scope.launch {
+            try {
+                nostr.waitUntilInitialized()
+                getCachedMetadata()
+            } catch (e: Exception) {
+                showError("Failed to load initial cache: ${e.message}")
+            }
         }
     }
 
-    private suspend fun runObserver() = coroutineScope {
-        launch {
-            nostr.profiles.metadataUpdates.collect { (pubkey, metadata) ->
-                updateMetadata(pubkey, Profile(pubkey, metadata))
-            }
+    private suspend fun runObserver() {
+        nostr.profiles.metadataUpdates.collect { (pubkey, metadata) ->
+            updateMetadata(pubkey, Profile(pubkey, metadata))
         }
     }
 
@@ -57,37 +70,69 @@ class ProfileCache(
                 batch.add(nextKey)
             }
 
-            nostr.profiles.fetchMetadataBatch(batch.toList())
-        }
-    }
-
-    private suspend fun loadCacheMetadata() {
-        val cache = nostr.profiles.getAllCacheMetadata()
-
-        profilesMutex.withLock {
-            cache.forEach { (pubkey, metadata) ->
-                val profile = Profile(pubkey, metadata)
-                profiles.getOrPut(pubkey) { MutableStateFlow(null) }.value = profile
-                seenPublicKeys.add(pubkey)
+            try {
+                nostr.profiles.fetchMetadataBatch(batch.toList())
+            } catch (e: Exception) {
+                // Allow these keys to be requested again since the fetch failed
+                seenPublicKeys.update { it - batch }
+                println("Failed to fetch metadata batch: ${e.message}")
             }
         }
     }
 
-    private fun requestMetadata(pubkey: PublicKey) {
-        if (seenPublicKeys.add(pubkey)) {
-            metadataRequestChannel.trySend(pubkey)
+    private suspend fun getCachedMetadata() {
+        val cache = nostr.profiles.getAllCacheMetadata()
+        cache.forEach { (pubkey, metadata) ->
+            updateMetadata(pubkey, Profile(pubkey, metadata))
         }
     }
 
-    private suspend fun updateMetadata(pubkey: PublicKey, profile: Profile) {
-        profilesMutex.withLock {
-            profiles.getOrPut(pubkey) { MutableStateFlow(null) }.value = profile
+    private fun requestMetadata(pubkey: PublicKey) {
+        var added = false
+        seenPublicKeys.update { current ->
+            if (current.contains(pubkey)) {
+                added = false
+                current
+            } else {
+                added = true
+                current + pubkey
+            }
         }
+        if (added) metadataRequestChannel.trySend(pubkey)
+    }
+
+    private fun updateMetadata(pubkey: PublicKey, profile: Profile) {
+        profiles.update { current ->
+            val flow = current[pubkey] ?: MutableStateFlow<Profile?>(null)
+            flow.value = profile
+            if (current.containsKey(pubkey)) current else current + (pubkey to flow)
+        }
+        seenPublicKeys.update { it + pubkey }
     }
 
     fun getMetadata(pubkey: PublicKey): StateFlow<Profile?> {
-        val flow = profiles.getOrPut(pubkey) { MutableStateFlow(null) }
-        if (flow.value == null) requestMetadata(pubkey)
-        return flow.asStateFlow()
+        val currentMap = profiles.value
+        val existingFlow = currentMap[pubkey]
+
+        if (existingFlow != null) {
+            if (existingFlow.value == null) requestMetadata(pubkey)
+            return existingFlow.asStateFlow()
+        }
+
+        val newFlow = MutableStateFlow<Profile?>(null)
+        var resultFlow = newFlow
+
+        profiles.update { prev ->
+            if (prev.containsKey(pubkey)) {
+                resultFlow = prev[pubkey]!!
+                prev
+            } else {
+                prev + (pubkey to newFlow)
+            }
+        }
+
+        if (resultFlow.value == null) requestMetadata(pubkey)
+
+        return resultFlow.asStateFlow()
     }
 }