From 06448a60a56a02580aac515ebb938f1e51666193 Mon Sep 17 00:00:00 2001
From: Han Yin <han.yin@arm.com>
Date: Wed, 16 Apr 2025 14:28:49 -0700
Subject: [PATCH] UI: update UI ongoing model import's cancellation

---
 .../ui/screens/ModelsManagementScreen.kt      | 26 +++++++++++--------
 .../viewmodel/ModelsManagementViewModel.kt    | 26 ++++++++++++++++++-
 2 files changed, 40 insertions(+), 12 deletions(-)

diff --git a/examples/llama.android/app/src/main/java/com/example/llama/revamp/ui/screens/ModelsManagementScreen.kt b/examples/llama.android/app/src/main/java/com/example/llama/revamp/ui/screens/ModelsManagementScreen.kt
index eee215e54e..eccdf85a41 100644
--- a/examples/llama.android/app/src/main/java/com/example/llama/revamp/ui/screens/ModelsManagementScreen.kt
+++ b/examples/llama.android/app/src/main/java/com/example/llama/revamp/ui/screens/ModelsManagementScreen.kt
@@ -255,7 +255,7 @@ fun ModelsManagementScreen(
                         }
 
                         IconButton(
-                            onClick = {/* TODO-han.yin: implement filtering */ }
+                            onClick = {/* TODO-han.yin: implement filtering once metadata ready */ }
                         ) {
                             Icon(
                                 imageVector = Icons.Default.FilterAlt,
@@ -377,9 +377,7 @@ fun ModelsManagementScreen(
                                 state.uri, state.fileName, state.fileSize
                             )
                         },
-                        onCancel = {
-                            viewModel.resetManagementState()
-                        }
+                        onCancel = { viewModel.resetManagementState() }
                     )
                 }
 
@@ -388,11 +386,10 @@ fun ModelsManagementScreen(
                         fileName = state.fileName,
                         fileSize = state.fileSize,
                         isImporting = true,
+                        isCancelling = state.isCancelling,
                         progress = state.progress,
                         onConfirm = {},
-                        onCancel = {
-                            // TODO-han.yin:  viewModel.cancelImport()
-                        },
+                        onCancel = { viewModel.cancelOngoingLocalModelImport() },
                     )
                 }
 
@@ -467,6 +464,7 @@ fun ImportProgressDialog(
     fileName: String,
     fileSize: Long,
     isImporting: Boolean,
+    isCancelling: Boolean = false,
     progress: Float,
     onConfirm: () -> Unit,
     onCancel: () -> Unit
@@ -528,7 +526,7 @@ fun ImportProgressDialog(
 
                 Spacer(modifier = Modifier.height(16.dp))
 
-                // Informational text
+                // Additional information
                 if (isImporting) {
                     Text(
                         text = "This may take several minutes for large models",
@@ -536,6 +534,12 @@ fun ImportProgressDialog(
                         color = MaterialTheme.colorScheme.onSurfaceVariant,
                         textAlign = TextAlign.Center
                     )
+                } else if (isCancelling) {
+                    Text(
+                        text = "Cancelling ongoing import...",
+                        style = MaterialTheme.typography.bodyMedium,
+                        color = MaterialTheme.colorScheme.error
+                    )
                 }
             }
         },
@@ -546,9 +550,9 @@ fun ImportProgressDialog(
             }
         },
         dismissButton = {
-            if (!isImporting || progress < 0.7f) {
-                TextButton(onClick = onCancel, enabled = !isImporting) {
-                    Text(if (isImporting) "Cancel" else "Back")
+            if (!isImporting || (progress < 0.7f && !isCancelling)) {
+                TextButton(onClick = onCancel, enabled = !isCancelling) {
+                    Text("Cancel")
                 }
             }
         }
diff --git a/examples/llama.android/app/src/main/java/com/example/llama/revamp/viewmodel/ModelsManagementViewModel.kt b/examples/llama.android/app/src/main/java/com/example/llama/revamp/viewmodel/ModelsManagementViewModel.kt
index 7220080368..94f7d32bae 100644
--- a/examples/llama.android/app/src/main/java/com/example/llama/revamp/viewmodel/ModelsManagementViewModel.kt
+++ b/examples/llama.android/app/src/main/java/com/example/llama/revamp/viewmodel/ModelsManagementViewModel.kt
@@ -20,6 +20,7 @@ import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.stateIn
+import kotlinx.coroutines.flow.update
 import kotlinx.coroutines.launch
 import java.io.FileNotFoundException
 import javax.inject.Inject
@@ -113,6 +114,29 @@ class ModelsManagementViewModel @Inject constructor(
         }
     }
 
+    fun cancelOngoingLocalModelImport() = viewModelScope.launch {
+        viewModelScope.launch {
+            // First update UI to show we're attempting to cancel
+            _managementState.update { current ->
+                if (current is Importation.Importing) {
+                    current.copy(isCancelling = true)
+                } else {
+                    current
+                }
+            }
+
+            // Attempt to cancel
+            when (modelRepository.cancelImport()) {
+                null, true -> { _managementState.value = ModelManagementState.Idle }
+                false -> {
+                    _managementState.value = Importation.Error(
+                        message = "Failed to cancel import. Try again later."
+                    )
+                }
+            }
+        }
+    }
+
     // TODO-han.yin: Stub for now. Would need to investigate HuggingFace APIs
     fun importFromHuggingFace() {}
 
@@ -171,7 +195,7 @@ sealed class ModelManagementState {
 
     sealed class Importation : ModelManagementState() {
         data class Confirming(val uri: Uri, val fileName: String, val fileSize: Long) : Importation()
-        data class Importing(val progress: Float = 0f, val fileName: String, val fileSize: Long) : Importation()
+        data class Importing(val progress: Float = 0f, val fileName: String, val fileSize: Long, val isCancelling: Boolean = false) : Importation()
         data class Success(val model: ModelInfo) : Importation()
         data class Error(val message: String) : Importation()
     }