feat(chat): support /undo command with transcript reload and composer prefill (#1051)
This commit is contained in:
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()
|
||||
}
|
||||
@@ -896,6 +896,45 @@ class ChatViewModelTest {
|
||||
assertEquals("openai/gpt-4o", viewModel.uiState.value.currentSessionModel)
|
||||
}
|
||||
|
||||
@Test
|
||||
fun testUndoCommand_dispatchesAndPrefillsComposer() =
|
||||
runTest {
|
||||
val (viewModel, sessionId) = createViewModelWithSession()
|
||||
|
||||
every {
|
||||
HermesWsClient.request(
|
||||
WsMethods.COMMAND_DISPATCH,
|
||||
mapOf("name" to "undo", "arg" to "2", "session_id" to sessionId),
|
||||
any(),
|
||||
)
|
||||
} returns
|
||||
CompletableDeferred(
|
||||
mapOf(
|
||||
"type" to "prefill",
|
||||
"message" to "undone user message",
|
||||
"notice" to "↶ Undid 2 turns",
|
||||
),
|
||||
)
|
||||
|
||||
viewModel.sendMessage("/undo 2")
|
||||
advanceUntilIdle()
|
||||
|
||||
// Prefill text is staged in UI state
|
||||
assertEquals("undone user message", viewModel.uiState.value.pendingPrefillText)
|
||||
// Notice is displayed as system message with rewind target feedback
|
||||
assertTrue(
|
||||
viewModel.uiState.value.messages
|
||||
.any { it.content.contains("↶ Undid 2 turns") && it.content.contains("undone user message") },
|
||||
)
|
||||
// Verify slash usage was recorded
|
||||
assertEquals(1, viewModel.uiState.value.slashUsageCounts["/undo"])
|
||||
|
||||
// Consume prefill text
|
||||
viewModel.consumePendingPrefill()
|
||||
advanceUntilIdle()
|
||||
assertNull(viewModel.uiState.value.pendingPrefillText)
|
||||
}
|
||||
|
||||
// ── Connection / init tests ──────────────────────────────────────────────
|
||||
|
||||
@Test
|
||||
|
||||
@@ -82,6 +82,14 @@ class SlashCommandDispatcherTest {
|
||||
assertEquals(SlashResult.SideQuestion(""), dispatcher.dispatch("/btw "))
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `undo routes to Undo`() {
|
||||
assertEquals(SlashResult.Undo(""), dispatcher.dispatch("/undo"))
|
||||
assertEquals(SlashResult.Undo("2"), dispatcher.dispatch("/undo 2"))
|
||||
assertEquals(SlashResult.Undo("3"), dispatcher.dispatch("/UNDO 3"))
|
||||
assertEquals(SlashResult.Undo(""), dispatcher.dispatch("/Undo"))
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `NEW uppercase still routes to NewSession`() {
|
||||
// Dispatcher lower-cases before matching, so case must not matter.
|
||||
|
||||
@@ -27,6 +27,10 @@ class FakeChatMessageDao : ChatMessageDao {
|
||||
messageList.forEach { messages[it.id] = it }
|
||||
}
|
||||
|
||||
override suspend fun deleteMessagesForSession(sessionId: String) {
|
||||
messages.values.removeAll { it.sessionId == sessionId }
|
||||
}
|
||||
|
||||
/** Direct access for test setup — bypasses the suspend modifier. */
|
||||
fun addMessageDirect(message: ChatMessageEntity) {
|
||||
messages[message.id] = message
|
||||
|
||||
Reference in new issue
Block a user