diff --git a/agent/app/src/main/java/com/example/appfunctions/agent/ui/screens/agentdemo/AgentDemoScreen.kt b/agent/app/src/main/java/com/example/appfunctions/agent/ui/screens/agentdemo/AgentDemoScreen.kt
index 327c8cf..c8d340c 100644
--- a/agent/app/src/main/java/com/example/appfunctions/agent/ui/screens/agentdemo/AgentDemoScreen.kt
+++ b/agent/app/src/main/java/com/example/appfunctions/agent/ui/screens/agentdemo/AgentDemoScreen.kt
@@ -134,8 +134,13 @@ import kotlinx.coroutines.launch
 @Composable
 fun AgentDemoScreen(viewModel: AgentDemoViewModel = hiltViewModel()) {
     val uiState by viewModel.uiState.collectAsStateWithLifecycle()
+    val messageText by viewModel.messageText.collectAsStateWithLifecycle()
 
-    AgentDemoContent(uiState = uiState, onEvent = viewModel::onEvent)
+    AgentDemoContent(
+        uiState = uiState,
+        onEvent = viewModel::onEvent,
+        messageText = messageText,
+    )
 }
 
 @OptIn(ExperimentalMaterial3Api::class)
@@ -143,6 +148,7 @@ fun AgentDemoScreen(viewModel: AgentDemoViewModel = hiltViewModel()) {
 fun AgentDemoContent(
     uiState: AgentUiState,
     onEvent: (AgentUiEvent) -> Unit,
+    messageText: String = "",
     initialSidePanelVisible: Boolean = false,
 ) {
     val context = LocalContext.current
@@ -172,6 +178,7 @@ fun AgentDemoContent(
                         drawerState = drawerState,
                         scope = scope,
                         packageManager = packageManager,
+                        messageText = messageText,
                         initialSidePanelVisible = initialSidePanelVisible,
                     )
                 }
@@ -222,9 +229,12 @@ fun AgentDemoLoadedScreen(
     drawerState: DrawerState,
     scope: CoroutineScope,
     packageManager: PackageManager,
+    messageText: String = "",
     initialSidePanelVisible: Boolean = false,
 ) {
-    var messageText by remember { mutableStateOf(TextFieldValue("")) }
+    var textFieldValue by remember {
+        mutableStateOf(TextFieldValue(messageText, TextRange(messageText.length)))
+    }
     var isSidePanelVisible by remember { mutableStateOf(initialSidePanelVisible) }
     var selectedAppPackageName by remember { mutableStateOf<String?>(null) }
 
@@ -337,15 +347,15 @@ fun AgentDemoLoadedScreen(
                 }
 
                 val sendMessage = {
-                    val textStr = messageText.text
+                    val textStr = textFieldValue.text
                     if (textStr.isNotBlank() && uiState.status == AgentStatus.Idle) {
                         onEvent(AgentUiEvent.OnSendMessage(textStr, selectedAppPackageName))
-                        messageText = TextFieldValue("")
+                        textFieldValue = TextFieldValue("")
                         selectedAppPackageName = null
                     }
                 }
 
-                val textStr = messageText.text
+                val textStr = textFieldValue.text
                 val lastAtIndex = textStr.lastIndexOf('@')
                 val showAutocomplete =
                     lastAtIndex >= 0 &&
@@ -401,9 +411,10 @@ fun AgentDemoLoadedScreen(
                 // Input area
                 Box(modifier = Modifier.fillMaxWidth()) {
                     OutlinedTextField(
-                        value = messageText,
+                        value = textFieldValue,
                         onValueChange = { newValue ->
-                            messageText = newValue
+                            textFieldValue = newValue
+                            onEvent(AgentUiEvent.OnMessageTextChanged(newValue.text))
                             val currentText = newValue.text
                             if (selectedAppPackageName != null && appMentionRegex != null) {
                                 if (!appMentionRegex.containsMatchIn(currentText)) {
@@ -441,7 +452,7 @@ fun AgentDemoLoadedScreen(
                             IconButton(
                                 onClick = sendMessage,
                                 enabled =
-                                    messageText.text.isNotBlank() &&
+                                    textFieldValue.text.isNotBlank() &&
                                         uiState.status == AgentStatus.Idle,
                             ) {
                                 Icon(
@@ -473,8 +484,8 @@ fun AgentDemoLoadedScreen(
                                         DropdownMenuItem(
                                             text = { Text(app.label) },
                                             onClick = {
-                                                val currentText = messageText.text
-                                                val selectionStart = messageText.selection.start
+                                                val currentText = textFieldValue.text
+                                                val selectionStart = textFieldValue.selection.start
                                                 val textBeforeCursor =
                                                     currentText.take(
                                                         selectionStart,
@@ -494,7 +505,7 @@ fun AgentDemoLoadedScreen(
                                                         "$textBeforeMention@${app.label} $textAfterCursor"
                                                     val newCursorPosition =
                                                         mentionIndex + app.label.length + 2
-                                                    messageText =
+                                                    textFieldValue =
                                                         TextFieldValue(
                                                             text = newText,
                                                             selection =
@@ -502,6 +513,9 @@ fun AgentDemoLoadedScreen(
                                                                     newCursorPosition,
                                                                 ),
                                                         )
+                                                    onEvent(
+                                                        AgentUiEvent.OnMessageTextChanged(newText),
+                                                    )
                                                     selectedAppPackageName = app.packageName
                                                 }
                                             },
@@ -920,12 +934,9 @@ fun ChatHistorySidePanel(
 
                 Surface(
                     modifier =
-                        Modifier
-                            .fillMaxWidth()
-                            .padding(vertical = 4.dp)
-                            .clickable {
-                                onEvent(AgentUiEvent.OnThreadSelected(thread.threadId))
-                            },
+                        Modifier.fillMaxWidth().padding(vertical = 4.dp).clickable {
+                            onEvent(AgentUiEvent.OnThreadSelected(thread.threadId))
+                        },
                     shape = MaterialTheme.shapes.medium,
                     color = backgroundColor,
                     contentColor = textColor,
diff --git a/agent/app/src/main/java/com/example/appfunctions/agent/ui/screens/agentdemo/AgentDemoViewModel.kt b/agent/app/src/main/java/com/example/appfunctions/agent/ui/screens/agentdemo/AgentDemoViewModel.kt
index 2f92660..8595e74 100644
--- a/agent/app/src/main/java/com/example/appfunctions/agent/ui/screens/agentdemo/AgentDemoViewModel.kt
+++ b/agent/app/src/main/java/com/example/appfunctions/agent/ui/screens/agentdemo/AgentDemoViewModel.kt
@@ -70,6 +70,9 @@ class AgentDemoViewModel
 
         private val installedApps = MutableStateFlow<List<AppInfo>>(emptyList())
 
+        val messageText: StateFlow<String> =
+            savedStateHandle.getStateFlow(KEY_MESSAGE_TEXT, "")
+
         private var observationJob: Job? = null
 
         init {
@@ -173,6 +176,7 @@ class AgentDemoViewModel
             val currentState = _uiState.value
             when (event) {
                 is AgentUiEvent.OnSendMessage -> {
+                    savedStateHandle[KEY_MESSAGE_TEXT] = ""
                     if (currentState is AgentUiState.Loaded) {
                         viewModelScope.launch {
                             sendMessageUseCase(
@@ -185,6 +189,9 @@ class AgentDemoViewModel
                         }
                     }
                 }
+                is AgentUiEvent.OnMessageTextChanged -> {
+                    savedStateHandle[KEY_MESSAGE_TEXT] = event.text
+                }
                 is AgentUiEvent.OnModelSelected -> {
                     if (currentState is AgentUiState.Loaded) {
                         viewModelScope.launch {
@@ -216,6 +223,8 @@ class AgentDemoViewModel
         }
     }
 
+private const val KEY_MESSAGE_TEXT = "messageText"
+
 private data class ThreadConfig(
     val threads: List<ThreadEntity>,
     val provider: LlmProviderName,
diff --git a/agent/app/src/main/java/com/example/appfunctions/agent/ui/screens/agentdemo/AgentUiState.kt b/agent/app/src/main/java/com/example/appfunctions/agent/ui/screens/agentdemo/AgentUiState.kt
index 4798f79..10ca0f4 100644
--- a/agent/app/src/main/java/com/example/appfunctions/agent/ui/screens/agentdemo/AgentUiState.kt
+++ b/agent/app/src/main/java/com/example/appfunctions/agent/ui/screens/agentdemo/AgentUiState.kt
@@ -49,4 +49,6 @@ sealed class AgentUiEvent {
     data class OnThreadSelected(val threadId: String) : AgentUiEvent()
 
     data class OnConfirmAction(val pendingIntentId: String) : AgentUiEvent()
+
+    data class OnMessageTextChanged(val text: String) : AgentUiEvent()
 }