feat(chat): support /undo command with transcript reload and composer prefill (#1051)

This commit is contained in:
M57 authored and GitHub committed 2026-09-03 19:04:43 +03:00
1 parent efe6fe81cb
commit b3439d87c9
8 files changed
+159

No files matched your search

@@ -18,4 +18,7 @@ interface ChatMessageDao {
@Insert(onConflict = OnConflictStrategy.REPLACE)
suspend fun upsertAll(messages: List<ChatMessageEntity>)
@Query("DELETE FROM chat_messages WHERE session_id = :sessionId")
suspend fun deleteMessagesForSession(sessionId: String)
}
@@ -33,4 +33,9 @@ open class ChatPersistenceRepository(
/** Load cached messages for a session from Room. */
suspend fun loadMessages(sessionId: String): List<ChatMessage> =
dao.getMessagesForSession(sessionId).map { it.toUiModel() }
/** Clear all cached messages for a session (e.g. after /undo rewind). */
suspend fun clearMessagesForSession(sessionId: String) {
dao.deleteMessagesForSession(sessionId)
}
}
@@ -290,6 +290,14 @@ fun ChatScreen(
var inputFieldValue by rememberSaveable(stateSaver = TextFieldValue.Saver) {
mutableStateOf(TextFieldValue(""))
}
LaunchedEffect(state.pendingPrefillText) {
val prefill = state.pendingPrefillText
if (prefill != null) {
inputFieldValue = ChatInputPolicy.commandFieldValue(prefill)
viewModel.consumePendingPrefill()
}
}
var isListening by rememberSaveable { mutableStateOf(false) }
var lastAnimatedMessageId by rememberSaveable { mutableStateOf<String?>(null) }
var showReloginDialog by rememberSaveable { mutableStateOf(false) }
@@ -363,6 +363,8 @@ data class ChatUiState(
// Session resume recovery (desktop parity: bounded auto-retry + error UI)
val resumeError: String? = null,
val isResumeRetrying: Boolean = false,
/** Text staged to prefill the composer (e.g. from /undo). */
val pendingPrefillText: String? = null,
) {
/** Convenience — derived from [connectionStatus]. */
val isConnected: Boolean get() = connectionStatus == ConnectionStatus.CONNECTED
@@ -1270,6 +1272,12 @@ class ChatViewModel(
handleSlashCommand(target)
}
"prefill" -> {
val message = map["message"] as? String ?: ""
val notice = map["notice"] as? String ?: ""
handlePrefillResult(message, notice)
}
else -> {
val output = map["output"] as? String ?: map.toString()
addAssistantMessage(output)
@@ -1611,6 +1619,11 @@ class ChatViewModel(
}
}
if (result is SlashResult.Undo) {
handleUndoCommand(result.count)
return
}
// Block desktop/CLI-only + TUI-only commands that don't function on
// mobile (issue #576, deliverable #3). These are also hidden from the
// suggestion menu, but a user can still type one — intercept it here
@@ -1682,6 +1695,10 @@ class ChatViewModel(
handleSideQuestionCommand(result.question)
}
is SlashResult.Undo -> {
handleUndoCommand(result.count)
}
is SlashResult.RpcDispatch -> {
dispatchViaRpc(command)
}
@@ -1996,6 +2013,67 @@ class ChatViewModel(
}
}
private fun handleUndoCommand(count: String) {
val sessionId = runtimeSessionId
if (sessionId == null) {
addAssistantMessage("No active session to undo.")
return
}
viewModelScope.launch {
slashUsageStore.recordUse("/undo")
}
viewModelScope.launch(ioDispatcher) {
try {
val result =
wsClient
.request(
WsMethods.COMMAND_DISPATCH,
mapOf("name" to "undo", "arg" to count, "session_id" to sessionId),
).await()
handleDispatchResult(result)
} catch (e: HermesWsClient.HermesRpcException) {
addAssistantMessage(e.message ?: "Failed to undo.")
} catch (e: Exception) {
addAssistantMessage(e.message ?: "Failed to undo.")
}
}
}
private fun handlePrefillResult(
message: String,
notice: String,
) {
val storageSessionId = _uiState.value.currentSessionId
if (message.isNotBlank()) {
_uiState.update { it.copy(pendingPrefillText = message) }
}
val feedback =
when {
notice.isNotBlank() && message.isNotBlank() -> "$notice (Rewound to: \"$message\")"
notice.isNotBlank() -> notice
message.isNotBlank() -> "↶ Rewound to: \"$message\""
else -> "↶ Rewound"
}
if (storageSessionId != null) {
val generation = sessionGeneration
viewModelScope.launch {
withContext(ioDispatcher) {
repo.clearMessagesForSession(storageSessionId)
}
_uiState.update { it.copy(messages = emptyList()) }
loadSessionMessages(storageSessionId, generation)
addSystemMessage(feedback, persist = true)
fetchContextUsage()
}
} else {
addSystemMessage(feedback, persist = false)
}
}
fun consumePendingPrefill() {
_uiState.update { it.copy(pendingPrefillText = null) }
}
// ── Session management ───────────────────────────────────────────────
fun interruptSession() {
@@ -2659,6 +2737,7 @@ class ChatViewModel(
todos = emptyList(),
resumeError = null,
isResumeRetrying = false,
pendingPrefillText = null,
)
}
return generation
@@ -66,6 +66,11 @@ class SlashCommandDispatcher {
SlashResult.SideQuestion(question = arg)
}
"/undo" -> {
val arg = command.split(" ", limit = 2).getOrElse(1) { "" }.trim()
SlashResult.Undo(count = arg)
}
else -> {
SlashResult.RpcDispatch
}
@@ -143,4 +148,12 @@ sealed class SlashResult {
data class SideQuestion(
val question: String,
) : SlashResult()
/**
* Rewind conversation turns via the `undo` command.dispatch RPC.
* Rewinds server transcript, reloads active messages, and prefills the composer.
*/
data class Undo(
val count: String,
) : SlashResult()
}