diff --git a/plugins/AI-Core/build.gradle.kts b/plugins/AI-Core/build.gradle.kts index fa67f6cf..fa98f596 100644 --- a/plugins/AI-Core/build.gradle.kts +++ b/plugins/AI-Core/build.gradle.kts @@ -85,6 +85,8 @@ dependencies { testImplementation(files("../../libs/plugin-api.jar")) // plugin-api's prompt loader parses YAML with the host's copy; JVM tests need their own, same version testImplementation("org.snakeyaml:snakeyaml-engine:2.10") + // plugin-api's PromptConfigStore logs through the host's slf4j; JVM tests need their own, same version + testRuntimeOnly("org.slf4j:slf4j-api:2.0.12") testImplementation("junit:junit:4.13.2") testImplementation("io.mockk:mockk:1.13.8") testImplementation("org.json:json:20240303") diff --git a/plugins/AI-Core/src/main/assets/docs/index.html b/plugins/AI-Core/src/main/assets/docs/index.html index 80b012b6..54b96536 100644 --- a/plugins/AI-Core/src/main/assets/docs/index.html +++ b/plugins/AI-Core/src/main/assets/docs/index.html @@ -262,6 +262,43 @@

Reading the conversation

text. It appears under a reply once the agent has finished writing it. +

Editing a message you sent

+

Tap the pencil under any of your messages to fix a typo, add a detail you + left out, or rephrase a request the agent misread. The message goes back into + the message box with the files that were attached to it, under an + Editing message banner.

+ + +

Versions of a chat

+

A message you edited from further up shows ‹ 2/2 › + underneath. Tap the arrows to switch between its versions; the chat below it + changes to what was said in that version.

+ +

Keeping several chats

Chats belong to the project they were started in. Opening a project shows that project's chats and no one else's, and you can keep as many as you want — diff --git a/plugins/AI-Core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/adapters/ChatAdapter.kt b/plugins/AI-Core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/adapters/ChatAdapter.kt index 2a311873..06f35cfc 100644 --- a/plugins/AI-Core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/adapters/ChatAdapter.kt +++ b/plugins/AI-Core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/adapters/ChatAdapter.kt @@ -1,5 +1,7 @@ package com.itsaky.androidide.plugins.aicore.adapters +import android.animation.ValueAnimator +import android.graphics.drawable.GradientDrawable import android.view.LayoutInflater import android.view.View import android.view.ViewGroup @@ -19,6 +21,7 @@ import com.itsaky.androidide.plugins.aicore.models.ChatMessage import com.itsaky.androidide.plugins.aicore.models.MessageStatus import com.itsaky.androidide.plugins.aicore.models.Sender import com.itsaky.androidide.plugins.aicore.plugin.AiCorePlugin +import com.itsaky.androidide.plugins.aicore.viewmodel.ChatBranches import io.noties.markwon.Markwon import java.text.DecimalFormat import java.text.SimpleDateFormat @@ -48,6 +51,45 @@ class ChatAdapter( private val decimalSecondsFormatter = DecimalFormat("0.0") private val expandedMessageIds = mutableSetOf() private val animatingHolders = mutableSetOf() + private val pulsingHolders = mutableSetOf() + + /** + * Whether user messages offer Edit and version switching, from ChatViewModel.canChangePrompts; + * false while a run is in flight. + */ + var canChangePrompts: Boolean = false + set(value) { + if (field == value) return + field = value + notifyUserRowsChanged { true } + } + + /** Every user message with other versions, and which of them it is; see ChatBranches. */ + internal var promptVersions: Map = emptyMap() + set(value) { + val previous = field + if (previous == value) return + field = value + notifyUserRowsChanged { previous[it] != value[it] } + } + + /** The user message loaded into the composer for editing, or null; ChatFragment owns the state. */ + var editingMessageId: String? = null + set(value) { + val previous = field + if (previous == value) return + field = value + notifyUserRowsChanged { it == previous || it == value } + } + + /** Rebinds the actions of the user rows [changed] picks; a payload keeps their text as it is. */ + private inline fun notifyUserRowsChanged(changed: (messageId: String) -> Boolean) { + currentList.forEachIndexed { position, message -> + if (message.sender == Sender.USER && changed(message.id)) { + notifyItemChanged(position, PromptActionsPayload) + } + } + } companion object { private const val VIEW_TYPE_DEFAULT = 0 @@ -56,7 +98,13 @@ class ChatAdapter( const val ACTION_RETRY = "retry" const val ACTION_COPY = "copy" + const val ACTION_EDIT = "edit" + const val ACTION_PREVIOUS_VERSION = "previous_version" + const val ACTION_NEXT_VERSION = "next_version" const val ACTION_OPEN_SETTINGS = "open_settings" + + /** A version arrow with nowhere to go, or held back while the agent works. */ + private const val DISABLED_ALPHA = 0.38f } sealed class MessageViewHolder(view: View) : RecyclerView.ViewHolder(view) @@ -72,8 +120,15 @@ class ChatAdapter( val btnRetry: Button = view.findViewById(R.id.btn_retry) /** Fold toggle; only the user bubble layout has one. */ val btnToggleExpand: ImageButton? = view.findViewById(R.id.btn_toggle_expand) + /** The bordered bubble; only the user bubble layout has one. */ + val userBubble: View? = view.findViewById(R.id.user_bubble) val messageActions: LinearLayout = view.findViewById(R.id.message_actions) val btnCopyMessage: ImageButton = view.findViewById(R.id.btn_copy_message) + val btnEditMessage: ImageButton = view.findViewById(R.id.btn_edit_message) + val versionNav: LinearLayout = view.findViewById(R.id.version_nav) + val btnPreviousVersion: ImageButton = view.findViewById(R.id.btn_previous_version) + val versionLabel: TextView = view.findViewById(R.id.version_label) + val btnNextVersion: ImageButton = view.findViewById(R.id.btn_next_version) /** * Queued next step of the "..." animation, or null when it isn't running. Retained so @@ -81,6 +136,9 @@ class ChatAdapter( * would keep this holder, its views and their Context reachable after the row is gone. */ var generatingDotsStep: Runnable? = null + + /** Pulses [userBubble]'s border while its prompt is being edited; null otherwise. */ + var editPulse: ValueAnimator? = null } class SystemMessageViewHolder(view: View) : MessageViewHolder(view) { @@ -151,6 +209,8 @@ class ChatAdapter( if (payloads.isEmpty()) { // No payload, do full bind onBindViewHolder(holder, position) + } else if (payloads.all { it === PromptActionsPayload }) { + if (holder is DefaultMessageViewHolder) updateMessageActions(holder, getItem(position)) } else { // Handle payload update val payload = payloads[0] @@ -265,11 +325,23 @@ class ChatAdapter( /** Wired once per holder: a streamed reply grows via payloads, so the tap reads the current item. */ private fun wireMessageActions(holder: DefaultMessageViewHolder) { - holder.btnCopyMessage.setOnClickListener { + holder.btnCopyMessage.sendsAction(holder, ACTION_COPY) + wireTooltip(holder.btnCopyMessage, AiCorePlugin.TOOLTIP_TAG_MESSAGE_COPY) + holder.btnEditMessage.sendsAction(holder, ACTION_EDIT) + wireTooltip(holder.btnEditMessage, AiCorePlugin.TOOLTIP_TAG_MESSAGE_EDIT) + holder.btnPreviousVersion.sendsAction(holder, ACTION_PREVIOUS_VERSION) + holder.btnNextVersion.sendsAction(holder, ACTION_NEXT_VERSION) + wireTooltip(holder.btnPreviousVersion, AiCorePlugin.TOOLTIP_TAG_MESSAGE_VERSIONS) + wireTooltip(holder.btnNextVersion, AiCorePlugin.TOOLTIP_TAG_MESSAGE_VERSIONS) + wireTooltip(holder.versionLabel, AiCorePlugin.TOOLTIP_TAG_MESSAGE_VERSIONS) + } + + /** Taps report [action] on the item [holder] shows now, which a streamed reply may have grown. */ + private fun View.sendsAction(holder: RecyclerView.ViewHolder, action: String) { + setOnClickListener { val pos = holder.bindingAdapterPosition - if (pos != RecyclerView.NO_POSITION) onMessageAction(ACTION_COPY, getItem(pos)) + if (pos != RecyclerView.NO_POSITION) onMessageAction(action, getItem(pos)) } - wireTooltip(holder.btnCopyMessage, AiCorePlugin.TOOLTIP_TAG_MESSAGE_COPY) } /** @@ -283,6 +355,61 @@ class ChatAdapter( val copyable = message.sender == Sender.USER || message.sender == Sender.AGENT val show = copyable && message.status != MessageStatus.LOADING && !streaming holder.messageActions.visibility = if (show) View.VISIBLE else View.GONE + val isUser = message.sender == Sender.USER + holder.btnEditMessage.visibility = if (isUser && canChangePrompts) View.VISIBLE else View.GONE + // Disabled outright: the prompt is already in the composer, and Cancel carries the tooltip. + val editing = isUser && message.id == editingMessageId + holder.btnEditMessage.isEnabled = !editing + holder.btnEditMessage.alpha = if (editing) DISABLED_ALPHA else 1f + if (editing) startEditPulse(holder) else stopEditPulse(holder) + val version = promptVersions[message.id]?.takeIf { isUser } + holder.versionNav.visibility = if (version != null) View.VISIBLE else View.GONE + if (version != null) { + val context = holder.versionLabel.context + holder.versionLabel.text = + context.getString(R.string.message_version_label, version.index + 1, version.count) + holder.versionLabel.contentDescription = + context.getString(R.string.desc_message_version, version.index + 1, version.count) + // Dimmed rather than disabled: a disabled view drops long-press, and with it the tooltip. + val canGoOlder = canChangePrompts && version.index > 0 + val canGoNewer = canChangePrompts && version.index < version.count - 1 + holder.btnPreviousVersion.alpha = if (canGoOlder) 1f else DISABLED_ALPHA + holder.btnNextVersion.alpha = if (canGoNewer) 1f else DISABLED_ALPHA + } + } + + /** + * Fades the bubble's border between its usual grey and the editing blue, and back, until + * [stopEditPulse]. Leaves a pulse already running alone, so a rebind doesn't restart it. + */ + private fun startEditPulse(holder: DefaultMessageViewHolder) { + val bubble = holder.userBubble ?: return + if (holder.editPulse != null) return + bubble.setBackgroundResource(R.drawable.bg_user_message_editing) + // Mutated, so the pulse recolors this bubble only and not every user of the drawable. + val background = bubble.background.mutate() as GradientDrawable + val context = bubble.context + val strokeWidth = context.resources.getDimensionPixelSize(R.dimen.chat_user_bubble_editing_stroke_width) + holder.editPulse = ValueAnimator.ofArgb( + context.getColor(R.color.plugin_outline_variant), + context.getColor(R.color.plugin_editing), + ).apply { + duration = context.resources.getInteger(R.integer.chat_user_bubble_editing_pulse_ms).toLong() + repeatMode = ValueAnimator.REVERSE + repeatCount = ValueAnimator.INFINITE + addUpdateListener { background.setStroke(strokeWidth, it.animatedValue as Int) } + start() + } + pulsingHolders.add(holder) + } + + /** Cancels the pulse, if any, and gives the bubble its usual border back. */ + private fun stopEditPulse(holder: DefaultMessageViewHolder) { + val pulse = holder.editPulse ?: return + pulse.cancel() + holder.editPulse = null + pulsingHolders.remove(holder) + holder.userBubble?.setBackgroundResource(R.drawable.bg_user_message) } private fun bindExpandToggle(holder: DefaultMessageViewHolder, toggle: ImageButton, message: ChatMessage) { @@ -417,18 +544,31 @@ class ChatAdapter( } /** - * Stop every live "…" animation. Call from the host fragment's `onDestroyView`: + * Stop every live "…" animation and edit pulse. Call from the host fragment's `onDestroyView`: * a message still streaming when the tab closes never reaches a terminal status * and its holder is never recycled, so nothing else cancels its Runnable. */ fun stopAllAnimations() { animatingHolders.toList().forEach { hideGeneratingDots(it) } + pulsingHolders.toList().forEach { stopEditPulse(it) } + } + + // A row scrolled off but still cached keeps its pulse; it only ticks while on screen. + override fun onViewDetachedFromWindow(holder: RecyclerView.ViewHolder) { + super.onViewDetachedFromWindow(holder) + (holder as? DefaultMessageViewHolder)?.editPulse?.pause() + } + + override fun onViewAttachedToWindow(holder: RecyclerView.ViewHolder) { + super.onViewAttachedToWindow(holder) + (holder as? DefaultMessageViewHolder)?.editPulse?.resume() } override fun onViewRecycled(holder: RecyclerView.ViewHolder) { super.onViewRecycled(holder) if (holder is DefaultMessageViewHolder) { hideGeneratingDots(holder) + stopEditPulse(holder) // Here, not in bind: a rebind of the same row would flash its toggle off for a frame. holder.btnToggleExpand?.visibility = View.GONE } @@ -520,4 +660,7 @@ class ChatAdapter( // Payload for partial updates data class TextUpdatePayload(val text: String, val status: MessageStatus) + + /** Payload for a user row whose Edit or version controls changed; nothing else about it did. */ + object PromptActionsPayload } diff --git a/plugins/AI-Core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/fragments/ChatFragment.kt b/plugins/AI-Core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/fragments/ChatFragment.kt index 4b3d594f..a85ac546 100644 --- a/plugins/AI-Core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/fragments/ChatFragment.kt +++ b/plugins/AI-Core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/fragments/ChatFragment.kt @@ -35,6 +35,7 @@ import com.itsaky.androidide.plugins.aicore.logging.AgentTrace import com.itsaky.androidide.plugins.aicore.logging.LOG_PREFIX import com.itsaky.androidide.plugins.aicore.managers.ProjectKey import com.itsaky.androidide.plugins.aicore.models.AgentState +import com.itsaky.androidide.plugins.aicore.models.ChatMessage import com.itsaky.androidide.plugins.aicore.models.ChatTranscript import com.itsaky.androidide.plugins.aicore.models.SessionRow import com.itsaky.androidide.plugins.aicore.models.isRunning @@ -77,6 +78,12 @@ class ChatFragment : Fragment(), ApprovalDialogFragment.Host { /** Saved-state key for [pendingExportSessionId], which must outlive the picker's round trip. */ const val STATE_PENDING_EXPORT = "pending_export_session_id" + + /** Saved-state keys for an edit in progress; see [restoreEdit]. */ + const val STATE_EDITING_ID = "editing_message_id" + const val STATE_DRAFT_BEFORE_EDIT = "draft_before_edit" + const val STATE_FILES_BEFORE_EDIT = "files_before_edit" + const val STATE_EDIT_FILES = "edit_files" } private var _binding: FragmentChatBinding? = null @@ -126,6 +133,16 @@ class ChatFragment : Fragment(), ApprovalDialogFragment.Host { if (uri != null) readImport(uri) } + /** + * The sent prompt loaded back into the composer for editing, or null outside edit mode. Send + * then replaces that prompt instead of adding a new one. + */ + private var editingMessageId: String? = null + + /** What the composer held before Edit borrowed it, put back when the edit ends either way. */ + private var draftBeforeEdit: String = "" + private var filesBeforeEdit: List = emptyList() + /** The message list's layout-declared padding, before any cutout inset is added. */ private val basePadding = Rect() @@ -195,6 +212,11 @@ class ChatFragment : Fragment(), ApprovalDialogFragment.Host { if (::chatAdapter.isInitialized) { chatAdapter.stopAllAnimations() } + // The composer goes with the view, so the next prompt must carry the draft's files again. + if (editingMessageId != null) viewModel.setContextFiles(filesBeforeEdit) + editingMessageId = null + draftBeforeEdit = "" + filesBeforeEdit = emptyList() super.onDestroyView() // runInFlight=true here, followed by that run still reporting, is the tab-switch fix // working; the run being gone from the trace after it is the bug coming back. Guarded @@ -230,6 +252,36 @@ class ChatFragment : Fragment(), ApprovalDialogFragment.Host { composer?.saveState(outState) // The same holds while the export picker is up, which is exactly when the IDE is backgrounded. outState.putString(STATE_PENDING_EXPORT, pendingExportSessionId) + // A tab switch rebuilds this fragment and restores the edited text, so edit mode must follow. + val editingId = editingMessageId ?: return + outState.putString(STATE_EDITING_ID, editingId) + outState.putString(STATE_DRAFT_BEFORE_EDIT, draftBeforeEdit) + outState.putStringArrayList(STATE_FILES_BEFORE_EDIT, ArrayList(filesBeforeEdit.map { it.path })) + outState.putStringArrayList(STATE_EDIT_FILES, ArrayList(contextFiles.map { it.path })) + } + + override fun onViewStateRestored(savedInstanceState: Bundle?) { + super.onViewStateRestored(savedInstanceState) + restoreEdit(savedInstanceState ?: return) + } + + /** + * Re-enters the edit a tab switch interrupted, keeping the edited text the view just restored; + * ends it instead, draft back in place, when that prompt is gone or can no longer change. + */ + private fun restoreEdit(state: Bundle) { + val id = state.getString(STATE_EDITING_ID) ?: return + val edited = binding.promptInputEdittext.text?.toString().orEmpty() + val editFiles = state.getStringArrayList(STATE_EDIT_FILES).orEmpty().map(::File) + editingMessageId = id + draftBeforeEdit = state.getString(STATE_DRAFT_BEFORE_EDIT).orEmpty() + filesBeforeEdit = state.getStringArrayList(STATE_FILES_BEFORE_EDIT).orEmpty().map(::File) + val message = viewModel.messages.value.firstOrNull { it.id == id } + if (message == null || !viewModel.canChangePrompts.value) return endEdit() + beginEdit(message) + binding.promptInputEdittext.setText(edited) + binding.promptInputEdittext.setSelection(edited.length) + replaceContextFiles(editFiles) } override fun onCreate(savedInstanceState: Bundle?) { @@ -562,12 +614,14 @@ class ChatFragment : Fragment(), ApprovalDialogFragment.Host { binding.btnAddContext.setOnClickListener { showFilePicker() } + binding.btnCancelEdit.setOnClickListener { endEdit() } // Anchored on the bar, not promptInputEdittext: long-press there is the paste menu. wireTooltip(binding.btnAddContext, AiCorePlugin.TOOLTIP_TAG_CONTEXT_FILES) wireTooltip(binding.inputBarCard, AiCorePlugin.TOOLTIP_TAG_CHAT_INPUT) wireTooltip(binding.sendButton, AiCorePlugin.TOOLTIP_TAG_CHAT_SEND) wireTooltip(binding.backendStatusText, AiCorePlugin.TOOLTIP_TAG_SETTINGS_BACKEND) + wireTooltip(binding.btnCancelEdit, AiCorePlugin.TOOLTIP_TAG_MESSAGE_EDIT) } /** @@ -580,11 +634,79 @@ class ChatFragment : Fragment(), ApprovalDialogFragment.Host { private fun sendPrompt() { val message = binding.promptInputEdittext.text?.toString() ?: return if (message.isBlank()) return + editingMessageId?.let { id -> + sendEdit(id, message) + return + } if (!viewModel.sendMessage(message)) return composer?.hideKeyboard() binding.promptInputEdittext.text?.clear() } + /** + * Confirms an edit: the chat rewinds or forks at [messageId] and runs on [text]. A refused send + * stays in edit mode with the text in place, as a refused plain send keeps its prompt. + */ + private fun sendEdit(messageId: String, text: String) { + when (viewModel.editPrompt(messageId, text)) { + ChatViewModel.EditResult.STARTED -> { + composer?.hideKeyboard() + endEdit() + } + ChatViewModel.EditResult.BUSY -> showInfoSnackbar(getString(R.string.msg_edit_stop_first)) + ChatViewModel.EditResult.NOT_EDITABLE -> { + endEdit() + showInfoSnackbar(getString(R.string.msg_edit_unavailable)) + } + ChatViewModel.EditResult.REFUSED -> Unit + } + } + + /** + * Loads a sent prompt and the files that went with it into the composer. Nothing is discarded + * until Send; the draft the composer held is kept aside and comes back when the edit ends. + */ + private fun beginEdit(message: ChatMessage) { + val binding = _binding ?: return + if (!viewModel.canChangePrompts.value) { + showInfoSnackbar(getString(R.string.msg_edit_stop_first)) + return + } + // Moving the edit to another prompt keeps the draft from before the first one. + if (editingMessageId == null) { + draftBeforeEdit = binding.promptInputEdittext.text?.toString().orEmpty() + filesBeforeEdit = contextFiles.toList() + } + editingMessageId = message.id + chatAdapter.editingMessageId = message.id + AgentTrace.detail("UI", "edit started files=${message.contextFiles?.size ?: 0}") + // Files deleted since the prompt was sent are left out rather than attached as empty. + replaceContextFiles(message.contextFiles.orEmpty().map(::File).filter { it.isFile }) + binding.promptInputEdittext.setText(message.text) + binding.promptInputEdittext.setSelection(binding.promptInputEdittext.length()) + // Only an older prompt forks, so only there does Send keep the original. + binding.editBannerText.setText( + if (viewModel.editForks(message.id)) R.string.chat_editing_message_fork + else R.string.chat_editing_message + ) + binding.editBanner.isVisible = true + composer?.focusInput() + } + + /** Leaves edit mode, putting back the draft and files the composer held before it began. */ + private fun endEdit() { + if (editingMessageId == null) return + editingMessageId = null + if (::chatAdapter.isInitialized) chatAdapter.editingMessageId = null + val binding = _binding ?: return + binding.editBanner.isVisible = false + binding.promptInputEdittext.setText(draftBeforeEdit) + binding.promptInputEdittext.setSelection(binding.promptInputEdittext.length()) + replaceContextFiles(filesBeforeEdit) + draftBeforeEdit = "" + filesBeforeEdit = emptyList() + } + /** * Records the message list's own padding so the cutout inset is added to it rather than * replacing it, and stays correct however many inset passes the window makes. @@ -670,6 +792,7 @@ class ChatFragment : Fragment(), ApprovalDialogFragment.Host { launch { observeMessages() } launch { observeAgentState() } launch { observePendingApprovalRequest() } + launch { observeEditableMessage() } } } } @@ -700,6 +823,33 @@ class ChatFragment : Fragment(), ApprovalDialogFragment.Host { } } + /** + * Keeps Edit and the version arrows in step with whether the agent is idle, and leaves edit + * mode once the prompt being edited is gone — a chat switch, a Retry, another version shown. + */ + private suspend fun observeEditableMessage() { + combine( + viewModel.canChangePrompts, + viewModel.promptVersions, + viewModel.messages, + ::Triple, + ).collect { (canChange, versions, messages) -> + chatAdapter.canChangePrompts = canChange + chatAdapter.promptVersions = versions + val editing = editingMessageId ?: return@collect + if (messages.none { it.id == editing }) endEdit() + } + } + + /** Shows the older ([step] -1) or newer (1) version of a prompt, if the agent is idle. */ + private fun switchVersion(message: ChatMessage, step: Int) { + if (!viewModel.canChangePrompts.value) { + showInfoSnackbar(getString(R.string.msg_version_stop_first)) + return + } + viewModel.switchPromptVersion(message.id, step) + } + /** True while the newest message is fully visible, i.e. the user is not reading back. */ private fun RecyclerView.isAtBottom(): Boolean = !canScrollVertically(1) @@ -825,6 +975,15 @@ class ChatFragment : Fragment(), ApprovalDialogFragment.Host { contextFiles.forEach(::addChipForFile) } + /** Swaps every attached-file chip for [files], and hands the ViewModel the same list. */ + private fun replaceContextFiles(files: List) { + val binding = _binding ?: return + contextFiles.clear() + binding.contextChipGroup.removeAllViews() + updateContextChipVisibility() + addContextFiles(files) + } + private fun addContextFiles(files: List) { files.forEach { file -> if (!contextFiles.contains(file)) { @@ -858,9 +1017,11 @@ class ChatFragment : Fragment(), ApprovalDialogFragment.Host { binding.contextChipScroll.isVisible = binding.contextChipGroup.childCount > 0 } - private fun onMessageAction(action: String, message: com.itsaky.androidide.plugins.aicore.models.ChatMessage) { + private fun onMessageAction(action: String, message: ChatMessage) { when (action) { ChatAdapter.ACTION_RETRY -> { + // First, so the retried prompt carries the draft's files and not the edit's. + endEdit() // The prompt behind this row, not its text: the row may be a tool failure. viewModel.retryLastRun() } @@ -869,6 +1030,9 @@ class ChatFragment : Fragment(), ApprovalDialogFragment.Host { openSettingsFragment() } ChatAdapter.ACTION_COPY -> copyToClipboard(message.text) + ChatAdapter.ACTION_EDIT -> beginEdit(message) + ChatAdapter.ACTION_PREVIOUS_VERSION -> switchVersion(message, -1) + ChatAdapter.ACTION_NEXT_VERSION -> switchVersion(message, 1) } } diff --git a/plugins/AI-Core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/fragments/ComposerAutoHideController.kt b/plugins/AI-Core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/fragments/ComposerAutoHideController.kt index 522ff134..c9859805 100644 --- a/plugins/AI-Core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/fragments/ComposerAutoHideController.kt +++ b/plugins/AI-Core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/fragments/ComposerAutoHideController.kt @@ -117,6 +117,16 @@ internal class ComposerAutoHideController( }) } + /** Opens the composer, if folded away, and puts the caret and keyboard in its input field. */ + fun focusInput() { + val binding = _binding ?: return + setVisible(true) + binding.promptInputEdittext.requestFocus() + val imm = binding.root.context + .getSystemService(Context.INPUT_METHOD_SERVICE) as InputMethodManager + imm.showSoftInput(binding.promptInputEdittext, InputMethodManager.SHOW_IMPLICIT) + } + /** Dismisses the soft keyboard raised by the input field. */ fun hideKeyboard() { val binding = _binding ?: return diff --git a/plugins/AI-Core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/models/ChatMessage.kt b/plugins/AI-Core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/models/ChatMessage.kt index e476b458..f17d7e45 100644 --- a/plugins/AI-Core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/models/ChatMessage.kt +++ b/plugins/AI-Core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/models/ChatMessage.kt @@ -56,5 +56,15 @@ data class ChatMessage( * once the backend answers as ready, so a key saved afterwards leaves no stranded warning * behind. False for every message stored before this field existed. */ - val isSetupError: Boolean = false + val isSetupError: Boolean = false, + /** + * Absolute paths of the context files attached when this user prompt was sent, so Edit can put + * the same chips back. Null on every other message, and on a prompt stored before this existed. + */ + val contextFiles: List? = null, + /** + * The message this one follows, set only on messages kept in ChatSession.otherBranches; null + * there means it opens the chat. On the branch on screen, order alone says what follows what. + */ + val parentId: String? = null ) diff --git a/plugins/AI-Core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/models/ChatSession.kt b/plugins/AI-Core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/models/ChatSession.kt index c9f59ff3..3c039f16 100644 --- a/plugins/AI-Core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/models/ChatSession.kt +++ b/plugins/AI-Core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/models/ChatSession.kt @@ -22,6 +22,12 @@ import java.util.UUID * @property generatedTitle the title the selected backend wrote after the first reply, or null * until then (and for sessions stored before this field, for the same Gson reason). Never * overrides [name]; see ChatViewModel's title generation. + * @property otherBranches every message of the versions not on screen, each with its parentId; + * [messages] is always the version on screen. Null until an older prompt is first edited, and + * for sessions stored before this field, for the same Gson reason. + * @property selectedBranches for a message that has several versions following it, the id of the + * one last on screen, so switching back returns there; the chat's start is keyed by + * ChatBranches.ROOT. Null alongside [otherBranches]. */ data class ChatSession( val id: String = UUID.randomUUID().toString(), @@ -29,7 +35,9 @@ data class ChatSession( val messages: List = emptyList(), val projectKey: String? = null, val name: String? = null, - val generatedTitle: String? = null + val generatedTitle: String? = null, + val otherBranches: List? = null, + val selectedBranches: Map? = null ) { /** * What names this chat in the session list: the [name] the user gave it, otherwise the diff --git a/plugins/AI-Core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/plugin/AiCorePlugin.kt b/plugins/AI-Core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/plugin/AiCorePlugin.kt index 2f325786..e62cf109 100644 --- a/plugins/AI-Core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/plugin/AiCorePlugin.kt +++ b/plugins/AI-Core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/plugin/AiCorePlugin.kt @@ -83,6 +83,10 @@ class AiCorePlugin : IPlugin, UIExtension, DocumentationExtension, SettingsExten const val TOOLTIP_TAG_SYSTEM_LOG = "agent_system_log" const val TOOLTIP_TAG_USER_MESSAGE_EXPAND = "agent_user_message_expand" const val TOOLTIP_TAG_MESSAGE_COPY = "agent_message_copy" + /** Edit on a user message, and the composer's cancel-edit button. */ + const val TOOLTIP_TAG_MESSAGE_EDIT = "agent_message_edit" + /** The "< 2/3 >" version switcher on a user message whose older prompt was edited. */ + const val TOOLTIP_TAG_MESSAGE_VERSIONS = "agent_message_versions" // Tags for the interactive controls on the AI Settings screen (see AiSettingsFragment). const val TOOLTIP_TAG_SETTINGS_BACK = "ai_settings_back" @@ -691,7 +695,47 @@ class AiCorePlugin : IPlugin, UIExtension, DocumentationExtension, SettingsExten detail = """

Copies the whole message, including any lines folded out of view, so you can paste it into a file or send it again.

-

More actions for a message will appear beside this one.

+

On your own messages, Edit sits beside it.

+ """.trimIndent(), + buttons = listOf( + PluginTooltipButton(description = "AI Core Agent guide", uri = "index.html", order = 0) + ) + ), + PluginTooltipEntry( + tag = TOOLTIP_TAG_MESSAGE_EDIT, + summary = "Edit this message and run the agent again from there.", + detail = """ +

Edit puts the message back in the message box, with the + files that were attached to it. Change it and tap Send, or + tap × to cancel and leave the chat as it was.

+

On your latest message, sending replaces it: the original + and everything the agent did after it are removed.

+

On an earlier message, sending starts a new version of the + chat from that point. The original is kept, and + ‹ 2/2 › under the message switches between + them. The agent only sees the messages above the one you edited.

+

Changes the agent already made to your files are not + undone. Check them, or ask the agent to revert them.

+

Not available while the agent is working or waiting for your + approval: stop it first.

+ """.trimIndent(), + buttons = listOf( + PluginTooltipButton(description = "AI Core Agent guide", uri = "index.html", order = 0) + ) + ), + PluginTooltipEntry( + tag = TOOLTIP_TAG_MESSAGE_VERSIONS, + summary = "Switch between the versions of this message.", + detail = """ +

This message was edited, so the chat has more than one version + from here on. ‹ and › show the other + versions; 2/3 says which one you are looking at.

+

Each version is its own conversation: the agent remembers only + the version on screen, never what was said in the others.

+

Files are not switched. Changes the agent made to your + project in one version stay when you show another.

+

Not available while the agent is working or waiting for your + approval: stop it first.

""".trimIndent(), buttons = listOf( PluginTooltipButton(description = "AI Core Agent guide", uri = "index.html", order = 0) diff --git a/plugins/AI-Core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/viewmodel/ChatBranches.kt b/plugins/AI-Core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/viewmodel/ChatBranches.kt new file mode 100644 index 00000000..4004c53a --- /dev/null +++ b/plugins/AI-Core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/viewmodel/ChatBranches.kt @@ -0,0 +1,189 @@ +package com.itsaky.androidide.plugins.aicore.viewmodel + +import com.itsaky.androidide.plugins.aicore.models.ChatMessage +import com.itsaky.androidide.plugins.aicore.models.ChatSession +import com.itsaky.androidide.plugins.aicore.models.Sender + +/** + * The versions a chat grows when an older prompt is edited: the edited prompt starts a new branch + * from the messages above it, and the original is kept beside it rather than discarded. + * + * A session's `messages` is always the branch on screen; every other branch's messages sit in + * `otherBranches`, each naming the message it follows. Kept free of Android and of the ViewModel, + * so the tree rules are unit-testable on their own. + */ +internal object ChatBranches { + + /** The `selectedBranches` key for the chat's start, which has no message of its own. */ + const val ROOT = "" + + /** + * Where a prompt sits among its versions. + * + * @property index zero-based, oldest version first. + * @property count how many versions there are; always at least two. + */ + data class Position(val index: Int, val count: Int) + + /** Oldest first; the id breaks a same-millisecond tie, so no order depends on what is on screen. */ + private val versionOrder = compareBy({ it.timestamp }, { it.id }) + + /** + * Every prompt on screen that has other versions, and where it sits among them. + * + * @param session the current chat. + * @return message id to position; empty for a chat that was never forked. + */ + fun positions(session: ChatSession): Map { + val others = session.otherBranches.orEmpty() + if (others.isEmpty()) return emptyMap() + val path = session.messages + val result = HashMap() + path.forEachIndexed { index, message -> + if (message.sender != Sender.USER) return@forEachIndexed + val versions = versionsAt(path, others, index) + if (versions.size > 1) { + result[message.id] = Position(versions.indexOfFirst { it.id == message.id }, versions.size) + } + } + return result + } + + /** + * Moves [messageId] and everything after it into `otherBranches`, leaving the messages above + * it on screen for the edited prompt to follow. Nothing is lost: the original is now a version. + * + * @param session the current chat. + * @param messageId a user prompt on screen. + * @return the forked session, or null when [messageId] is not a prompt on screen. + */ + fun fork(session: ChatSession, messageId: String): ChatSession? { + val path = session.messages + val index = path.indexOfFirst { it.id == messageId } + if (index < 0 || path[index].sender != Sender.USER) return null + val (others, selected) = archive(session, index) + return session.copy( + messages = path.subList(0, index).toList(), + otherBranches = others, + selectedBranches = selected, + ) + } + + /** + * Puts version [targetId] of the prompt [messageId] on screen, followed by whatever was last + * on screen after it; the branch it replaces is kept for switching back. + * + * @param session the current chat. + * @param messageId a user prompt on screen. + * @param targetId another version of that prompt, from `otherBranches`. + * @return the switched session, or null when [targetId] is not a version of [messageId]. + */ + fun switchTo(session: ChatSession, messageId: String, targetId: String): ChatSession? { + val path = session.messages + val index = path.indexOfFirst { it.id == messageId } + if (index < 0 || path[index].sender != Sender.USER) return null + val parent = parentIdAt(path, index) + val target = session.otherBranches.orEmpty().firstOrNull { + it.id == targetId && it.sender == Sender.USER && it.parentId == parent + } ?: return null + val (archived, selected) = archive(session, index) + val tail = descend(target, archived, selected) + val tailIds = tail.mapTo(HashSet()) { it.id } + return session.copy( + // Order says what follows what on screen, so the stored link is dropped there. + messages = path.subList(0, index) + tail.map { it.copy(parentId = null) }, + otherBranches = archived.filterNot { it.id in tailIds }, + selectedBranches = selected + ((parent ?: ROOT) to target.id), + ) + } + + /** + * The version of [messageId] [step] places away from it, oldest first. + * + * @return that version's id, or null past either end or when the prompt has one version. + */ + fun versionId(session: ChatSession, messageId: String, step: Int): String? { + val path = session.messages + val index = path.indexOfFirst { it.id == messageId } + if (index < 0) return null + val versions = versionsAt(path, session.otherBranches.orEmpty(), index) + val target = versions.indexOfFirst { it.id == messageId } + step + return versions.getOrNull(target)?.id?.takeUnless { it == messageId } + } + + /** + * Drops the on-screen messages [doomed] accepts, re-pointing every stored version that followed + * one of them at the nearest message kept above it, so no version is left following nothing. + * + * @param session the current chat. + * @param doomed picks the on-screen messages to remove; stored versions are never removed. + * @return the session without them; [session] itself when nothing matched. + */ + fun removeFromScreen(session: ChatSession, doomed: (ChatMessage) -> Boolean): ChatSession { + val path = session.messages + if (path.none(doomed)) return session + // What each removed message is replaced by as a parent: the nearest survivor above it. + val successor = HashMap() + var keptAbove: String? = null + for (message in path) { + if (doomed(message)) successor[message.id] = keptAbove else keptAbove = message.id + } + fun repoint(id: String?): String? = if (id != null && id in successor) successor[id] else id + return session.copy( + messages = path.filterNot(doomed), + otherBranches = session.otherBranches?.map { it.copy(parentId = repoint(it.parentId)) }, + selectedBranches = session.selectedBranches?.mapKeys { (key, _) -> + if (key == ROOT) key else repoint(key) ?: ROOT + }, + ) + } + + /** The prompt at [index] and every stored version of it, oldest first. */ + private fun versionsAt( + path: List, + others: List, + index: Int, + ): List { + val parent = parentIdAt(path, index) + val siblings = others.filter { it.sender == Sender.USER && it.parentId == parent } + return (siblings + path[index]).sortedWith(versionOrder) + } + + /** What the message at [index] follows: the one above it, or null at the chat's start. */ + private fun parentIdAt(path: List, index: Int): String? = + if (index == 0) null else path[index - 1].id + + /** + * Links the branch from [index] down into `otherBranches`, and records it as the one last on + * screen at every step, so a later switch back lands on it and not on an older sibling. + */ + private fun archive(session: ChatSession, index: Int): Pair, Map> { + val path = session.messages + val selected = session.selectedBranches.orEmpty().toMutableMap() + val tail = (index until path.size).map { i -> + val parent = parentIdAt(path, i) + selected[parent ?: ROOT] = path[i].id + path[i].copy(parentId = parent) + } + return (session.otherBranches.orEmpty() + tail) to selected + } + + /** + * [start] and the branch below it that was last on screen, falling back on the newest where + * nothing was recorded. + */ + private fun descend( + start: ChatMessage, + others: List, + selected: Map, + ): List { + val children = others.groupBy { it.parentId } + val branch = mutableListOf(start) + while (true) { + val next = children[branch.last().id] ?: break + branch += next.firstOrNull { it.id == selected[branch.last().id] } + ?: next.maxWith(versionOrder) + } + return branch + } +} diff --git a/plugins/AI-Core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/viewmodel/ChatViewModel.kt b/plugins/AI-Core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/viewmodel/ChatViewModel.kt index 1afb84eb..5818478e 100644 --- a/plugins/AI-Core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/viewmodel/ChatViewModel.kt +++ b/plugins/AI-Core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/viewmodel/ChatViewModel.kt @@ -53,7 +53,6 @@ import com.itsaky.androidide.plugins.services.SharedServices import java.io.File import java.io.IOException import java.util.UUID -import java.util.concurrent.atomic.AtomicBoolean import java.util.concurrent.atomic.AtomicInteger import java.util.concurrent.atomic.AtomicLong import kotlinx.coroutines.CancellationException @@ -71,6 +70,7 @@ import kotlinx.coroutines.flow.SharingStarted import kotlinx.coroutines.flow.StateFlow import kotlinx.coroutines.flow.asStateFlow import kotlinx.coroutines.flow.combine +import kotlinx.coroutines.flow.distinctUntilChangedBy import kotlinx.coroutines.flow.map import kotlinx.coroutines.flow.stateIn import kotlinx.coroutines.future.await @@ -335,7 +335,7 @@ class ChatViewModel( private val generationEpoch = AtomicInteger(0) /** True while a generation is admitted and its coroutine has not yet unwound; gates re-entry. */ - private val isGenerating = AtomicBoolean(false) + private val isGenerating = MutableStateFlow(false) /** Whether the current run's most recent tool batch failed; reset per run. */ @Volatile @@ -381,9 +381,48 @@ class ChatViewModel( @Volatile private var historySizeBeforeLastRun = 0 + /** + * The user message the last run posted, so [editPrompt] knows [historySizeBeforeLastRun] + * describes that prompt and not an earlier one; null when the retry point is forgotten. + */ + @Volatile + private var lastRunUserMessageId: String? = null + /** The tool awaiting approval, straight from [approvalManager] — no polling in between. */ val pendingApprovalRequest: StateFlow = approvalManager.currentApprovalRequest + /** + * Whether the user's prompts offer Edit and version switching: false while a run is generating, + * executing a tool, or waiting on an approval, which counts as busy on its own. A stopped run + * stays busy until its coroutine unwinds, which a blocking tool can hold up for minutes. + */ + val canChangePrompts: StateFlow = + combine(isGenerating, _agentState, pendingApprovalRequest, ::isIdle) + .stateIn(viewModelScope, SharingStarted.Eagerly, false) + + /** Every prompt on screen with other versions, and which of them it is; see [ChatBranches]. */ + internal val promptVersions: StateFlow> = + combine(_sessions, _currentSessionId) { sessions, sessionId -> + sessions.firstOrNull { it.id == sessionId } + } + // A streamed token replaces the session but moves no message, so it needs no recount. + .distinctUntilChangedBy { session -> + session?.let { Triple(it.id, it.messages.map(ChatMessage::id), it.otherBranches) } + } + .map { session -> session?.let(ChatBranches::positions).orEmpty() } + .stateIn(viewModelScope, SharingStarted.Eagerly, emptyMap()) + + /** The one rule [canChangePrompts] and [isBusy] share: no run in flight, no approval open. */ + private fun isIdle(generating: Boolean, state: AgentState, approval: ApprovalRequest?): Boolean = + !generating && !state.isRunning && approval == null + + /** + * Whether a prompt may not be edited or switched right now. Read from the sources rather than + * [canChangePrompts], which a collector can see a beat late. + */ + private fun isBusy(): Boolean = + !isIdle(isGenerating.value, _agentState.value, pendingApprovalRequest.value) + private var _contextFiles = listOf() /** Files the user attached; read back by the fragment to rebuild its chips on re-attach. */ @@ -626,7 +665,9 @@ class ChatViewModel( if (remaining.size == _messages.value.size) return false _messages.value = remaining val session = currentSessionOrNull() ?: return true - replaceCurrentSessionMessages(session.messages.filterNot(doomed)) + // Through ChatBranches: a stored version may follow a removed message, and must not dangle. + val trimmed = ChatBranches.removeFromScreen(session, doomed) + _sessions.value = _sessions.value.map { if (it.id == session.id) trimmed else it } return true } @@ -890,10 +931,24 @@ class ChatViewModel( * which is what keeps the composer's text in place for an unconfigured backend. */ fun sendMessage(userMessage: String): Boolean { + val llmService = passPreflight(userMessage) ?: return false + // Reject re-entry while a generation is still in flight. + if (!isGenerating.compareAndSet(false, true)) return false + launchRun(llmService, userMessage, contextFiles) + return true + } + + /** + * The checks a prompt must pass before any run starts; a setup failure leaves a notice. + * + * @param userMessage the prompt about to be sent. + * @return the service to run against, or null when the prompt must not be sent. + */ + private fun passPreflight(userMessage: String): LlmInferenceService? { val llmService = getLlmService() if (llmService == null) { emitSystemError(str(R.string.error_llm_service_not_available)) - return false + return null } if (!_backendStatus.value.isAvailable) { @@ -909,19 +964,22 @@ class ChatViewModel( SelectedBackend.None -> str(R.string.backend_none_installed) } ) - return false - } - - if (userMessage.isBlank()) { - return false + return null } - // Reject re-entry while a generation is still in flight. - if (!isGenerating.compareAndSet(false, true)) { - return false - } + return llmService.takeUnless { userMessage.isBlank() } + } - AgentTrace.beginRun(currentBackendId, userMessage, contextFiles.size) + /** + * Starts the agent run for a prompt that passed [passPreflight], once the caller has claimed + * [isGenerating]. Main only: the prompt is on screen by the time this returns. + * + * @param llmService the service [passPreflight] returned. + * @param userMessage the prompt to send. + * @param runFiles the attachments, read by the caller before anything it changed could move them. + */ + private fun launchRun(llmService: LlmInferenceService, userMessage: String, runFiles: List) { + AgentTrace.beginRun(currentBackendId, userMessage, runFiles.size) // Reset per-run tool tracking. lastToolFailedThisRun = false activityMessageId = null @@ -929,32 +987,36 @@ class ChatViewModel( runToolLog.clear() runCodeReply = null runChangedProject = false + // Created here, on Main, so the retry point below names the prompt it belongs to. + val userChatMessage = ChatMessage( + id = UUID.randomUUID().toString(), + text = userMessage, + sender = Sender.USER, + status = MessageStatus.SENT, + contextFiles = runFiles.map { it.absolutePath }.takeIf { it.isNotEmpty() }, + ) // Where a Retry has to rewind to; read here, on Main, while no run can be appending. lastRunPrompt = userMessage + lastRunUserMessageId = userChatMessage.id historySizeBeforeLastRun = _history.value.size // Read once: prompt, grammar and executor must all describe the same tool set. val tools = agentTools val epoch = generationEpoch.incrementAndGet() - generationJob = viewModelScope.launch(Dispatchers.IO) { - // The session this run may title, and whether a title request took that job over. - var titleSessionId: String? = null + // The session this run may title; read before the message lands, so the header never + // shows it as the title first. + val titleSessionId = currentSessionOrNull()?.takeIf { needsTitle(it) }?.id + titleSessionId?.let { id -> _titlePending.value = _titlePending.value + id } + // Posted now, not from the run: a Stop before the run's first dispatch left a fork with no + // prompt at its fork point, and so no arrows to reach the version it had put aside. + _messages.value = _messages.value + userChatMessage + syncMessageToSession(userChatMessage) + setState(AgentState.Processing(str(R.string.msg_generating))) + // ATOMIC: a Stop before the first dispatch skipped the block, finally too, so isGenerating + // stayed set and every later send was refused; now the cancel lands at the first suspension. + generationJob = viewModelScope.launch(Dispatchers.IO, start = CoroutineStart.ATOMIC) { + // Whether a title request took over settling this run's title placeholder. var titleRequestStarted = false try { - // Add user message to the UI. - val userChatMessage = ChatMessage( - id = UUID.randomUUID().toString(), - text = userMessage, - sender = Sender.USER, - status = MessageStatus.SENT - ) - withContext(Dispatchers.Main) { - // Before the message lands, so the header never shows it as the title first. - titleSessionId = currentSessionOrNull()?.takeIf { needsTitle(it) }?.id - titleSessionId?.let { id -> _titlePending.value = _titlePending.value + id } - _messages.value = _messages.value + userChatMessage - syncMessageToSession(userChatMessage) - setState(AgentState.Processing(str(R.string.msg_generating))) - } // Queued behind a title still being written; the prompt shows as generating meanwhile. awaitTitleRequest() @@ -980,7 +1042,7 @@ class ChatViewModel( val messageWithContext = buildString { append(userMessage) - append(contextFilesPrompt.render(contextFiles)) + append(contextFilesPrompt.render(runFiles)) } val history = _history.value.toMutableList() history.add( @@ -992,7 +1054,7 @@ class ChatViewModel( // Code to judge, or a question about what is current: search before answering. val requiredTool = WebAccess.WEB_SEARCH_TOOL.takeIf { search -> toolDefinitions.any { it.name == search } && - VerificationPolicy.requiresWebCheck(userMessage, contextFiles.isNotEmpty()) + VerificationPolicy.requiresWebCheck(userMessage, runFiles.isNotEmpty()) } requiredTool?.let { AgentTrace.stage("VERIFY", "required=$it on the first turn") } var firstTurn = true @@ -1039,9 +1101,10 @@ class ChatViewModel( } } } finally { - // Persist history only if this run wasn't superseded (epoch bumped). - if (generationEpoch.get() == epoch) { - _history.value = history.toList() + // On Main, so a Stop, clear or chat switch can't bump the epoch between check and write. + withContext(NonCancellable + Dispatchers.Main) { + // Persist history only if this run wasn't superseded (epoch bumped). + if (generationEpoch.get() == epoch) _history.value = history.toList() } stopStateTimer() } @@ -1059,7 +1122,7 @@ class ChatViewModel( addSystemMessage(str(R.string.state_error, e.message), MessageStatus.ERROR) } finally { // Allow re-entry once the coroutine unwinds. - isGenerating.set(false) + isGenerating.value = false // The run can finish with the chat off screen, where nothing else writes. // NonCancellable so a Stop still saves what the run produced before it, and so the // activity line is closed rather than left reading as a tool still running. @@ -1071,7 +1134,6 @@ class ChatViewModel( } } } - return true } /** @@ -1261,7 +1323,7 @@ class ChatViewModel( return } // Checked before the rewind, which would otherwise pull the transcript out of a live run. - if (isGenerating.get()) { + if (isGenerating.value) { logDebug("retryLastRun: a run is already in flight") return } @@ -1269,6 +1331,135 @@ class ChatViewModel( sendMessage(prompt) } + /** How [editPrompt] ended, which decides whether the composer leaves edit mode. */ + enum class EditResult { + /** The chat was rewound or forked and the edited prompt is running. */ + STARTED, + /** A run is in flight or an approval is open; the user has to stop it first. */ + BUSY, + /** The message is no longer a prompt on screen. */ + NOT_EDITABLE, + /** The edited text was blank or the backend is not ready; nothing was changed. */ + REFUSED, + } + + /** + * Runs the agent on [newText] in place of the prompt [messageId]. The newest prompt is replaced + * outright, as if it had never been sent; an older one forks, keeping the original as a version + * (see [ChatBranches]). Either way the model sees only the messages above it. File edits a + * run already made are not undone. + * + * @param messageId a user prompt on screen. + * @param newText the edited prompt, sent with the context files attached right now. + * @return what happened; only [EditResult.STARTED] changed anything. + */ + fun editPrompt(messageId: String, newText: String): EditResult { + // Before the pre-flight, whose setup notice would otherwise land inside a live run. + if (isBusy()) return EditResult.BUSY + if (_messages.value.none { it.id == messageId && it.sender == Sender.USER }) { + return EditResult.NOT_EDITABLE + } + // Read before the cut: the composer leaves edit mode on it, handing back the draft's files. + val runFiles = contextFiles + // Checked before the rewind, so a refused send changes nothing. + val llmService = passPreflight(newText) ?: return EditResult.REFUSED + // Claimed before the cut too, so a chat is never cut for a run that then cannot start. + if (!isGenerating.compareAndSet(false, true)) return EditResult.BUSY + val cut = if (editForks(messageId)) forkAt(messageId) else rewind(messageId) + if (!cut) { + isGenerating.value = false + return EditResult.NOT_EDITABLE + } + launchRun(llmService, newText, runFiles) + return EditResult.STARTED + } + + /** + * Whether editing the prompt [messageId] forks the chat, keeping the original as a version, + * rather than replacing it: true for every prompt but the newest. + */ + fun editForks(messageId: String): Boolean = !PromptEdit.isLatestPrompt(_messages.value, messageId) + + /** + * Shows the version of prompt [messageId] [step] places away, and gives the model that branch + * alone: nothing said in another version reaches it. Main only. + * + * @param step -1 for the older version, 1 for the newer one. + * @return false, changing nothing, while busy or past either end of the versions. + */ + fun switchPromptVersion(messageId: String, step: Int): Boolean { + if (isBusy()) return false + val session = currentSessionOrNull() ?: return false + val targetId = ChatBranches.versionId(session, messageId, step) ?: return false + val switched = ChatBranches.switchTo(session, messageId, targetId) ?: return false + adoptBranch(switched, "VERSION", "step=$step") + return true + } + + /** + * Forks the chat at the older prompt [messageId]: it and everything after it become a stored + * version, and the messages above it are all the next run sees. Main only; the caller checks + * that no run is in flight. + * + * @return false, changing nothing, when [messageId] is not a prompt on screen. + */ + private fun forkAt(messageId: String): Boolean { + val session = currentSessionOrNull() ?: return false + val forked = ChatBranches.fork(session, messageId) ?: return false + adoptBranch(forked, "FORK", "keptMessages=${forked.messages.size}") + return true + } + + /** + * Puts [session], a branch change of the current chat, on screen and rebuilds the model's + * history from its messages alone, then writes it at once. + */ + private fun adoptBranch(session: ChatSession, stage: String, detail: String) { + _sessions.value = _sessions.value.map { if (it.id == session.id) session else it } + showTranscript(session.messages) + traceHistoryCut(stage, detail) + persistState() + } + + /** Traces a change to the model's history, with how much of it was kept and its last turn. */ + private fun traceHistoryCut(stage: String, detail: String) { + AgentTrace.stage( + stage, + "$detail historyKept=${_history.value.size}", + _history.value.lastOrNull()?.let { "${it.role}: ${AgentTrace.preview(it.content)}" }, + ) + } + + /** + * Cuts the conversation back to just before [messageId]: transcript, session and history all + * lose that prompt and every turn after it, and the cut session is written at once. Main only. + * + * @param messageId the newest user prompt. + * @return false, changing nothing, while busy or when [messageId] is not the newest prompt. + */ + internal fun rewindTo(messageId: String): Boolean = !isBusy() && rewind(messageId) + + /** [rewindTo] for a caller that has already checked, and claimed, that no run is in flight. */ + private fun rewind(messageId: String): Boolean { + val kept = PromptEdit.messagesBefore(_messages.value, messageId) ?: return false + val dropped = _messages.value.size - kept.size + _messages.value = kept + replaceCurrentSessionMessages(kept) + // The live history carries the tool turns a rebuild cannot; use it when it is this prompt's. + _history.value = if (messageId == lastRunUserMessageId) { + _history.value.take(historySizeBeforeLastRun) + } else { + rebuildHistoryFrom(kept) + } + forgetRetryPoint() + // A title written from the discarded first exchange would name a conversation that is gone. + if (kept.none { it.sender == Sender.USER }) resetGeneratedTitle() + traceHistoryCut("EDIT", "droppedMessages=$dropped") + // Written now rather than debounced, like a clear: the discarded branch must not come back. + persistState() + return true + } + /** * Runs one streaming model turn: creates an agent bubble, streams tokens into it and suspends * until completion, throwing on backend error. Sends [turns] structurally to the local backend; @@ -1671,26 +1862,39 @@ class ChatViewModel( _history.value = emptyList() // Without this the session keeps its messages and the cleared chat returns on the next sync. replaceCurrentSessionMessages(emptyList()) - // The title described the conversation just cleared; the next first reply writes a new one. - // Its reply would name the cleared chat, and its settle would end the next run's placeholder. + // Its other versions go too: a cleared chat has no prompt left for them to belong to. _currentSessionId.value?.let { sessionId -> - // Another chat's request is left to finish and settle its own placeholder. - titleRequest?.takeIf { it.sessionId == sessionId }?.let { - it.job.cancel() - titleRequest = null - } - titleRequested.remove(sessionId) - settleTitle(sessionId) _sessions.value = _sessions.value.map { - if (it.id == sessionId) it.copy(generatedTitle = null) else it + if (it.id == sessionId) it.copy(otherBranches = null, selectedBranches = null) else it } } + // The title described the conversation just cleared; the next first reply writes a new one. + resetGeneratedTitle() forgetRetryPoint() setState(AgentState.Idle) // Written now rather than debounced: a clear is deliberate and must survive a force-stop. persistState() } + /** + * Drops the current chat's generated title, and any request still writing one, so the next + * first reply names it afresh. A title the user typed is kept. Main only. + */ + private fun resetGeneratedTitle() { + val sessionId = _currentSessionId.value ?: return + // Its reply would name the discarded chat, and its settle would end the next run's placeholder. + // Another chat's request is left to finish and settle its own placeholder. + titleRequest?.takeIf { it.sessionId == sessionId }?.let { + it.job.cancel() + titleRequest = null + } + titleRequested.remove(sessionId) + settleTitle(sessionId) + _sessions.value = _sessions.value.map { + if (it.id == sessionId) it.copy(generatedTitle = null) else it + } + } + /** * Stops any in-flight run, so the transcript it is streaming into cannot be swapped out from * under it: without the epoch bump the stale run's callbacks keep writing, and its reply lands @@ -1746,23 +1950,27 @@ class ChatViewModel( * Makes [session] the live conversation: what the screen shows, what the model is given and * what the next message is appended to. * - * The one place those three move together. They were being set side by side at each of the - * three call sites, which is how restoring a transcript to the screen while handing the model - * an empty array (ADFA-5584) could be fixed in one of them and not the others. - * * @param session the conversation to make current; it must already be in [_sessions]. */ private fun adoptSession(session: ChatSession) { _currentSessionId.value = session.id - // Immutable snapshot, so a later mutation cannot reach collectors behind the StateFlow. - _messages.value = session.messages.toList() - // Emptying this is what had the model forget a conversation the user was looking at. - _history.value = rebuildHistoryFrom(session.messages) + showTranscript(session.messages) // Also dropped here: the readiness edge may have fired while another chat was current. if (_backendStatus.value.isAvailable) clearBackendSetupNotices() + persistState() + } + + /** + * Puts [messages] on screen and rebuilds the model's history from them. The one place those + * move together: set side by side at each call site, ADFA-5584 (screen restored, model handed + * an empty history) was fixed in one of them and not the others. + */ + private fun showTranscript(messages: List) { + // Immutable snapshot, so a later mutation cannot reach collectors behind the StateFlow. + _messages.value = messages.toList() + _history.value = rebuildHistoryFrom(messages) // The rewind point names a run this transcript does not have, and would truncate it. forgetRetryPoint() - persistState() } /** @@ -1899,6 +2107,7 @@ class ChatViewModel( */ private fun forgetRetryPoint() { lastRunPrompt = null + lastRunUserMessageId = null historySizeBeforeLastRun = 0 } @@ -2060,13 +2269,11 @@ class ChatViewModel( * partial ones [MessageStatus.COMPLETED]. Called on Stop. */ private fun finalizeInProgressMessages() { - val finalized = _messages.value.mapNotNull { msg -> - if (msg.sender == Sender.AGENT && msg.durationMs == null) { - if (msg.text.isBlank()) null - else msg.copy(status = MessageStatus.COMPLETED, durationMs = 0L) - } else { - msg - } + val unfinished = { msg: ChatMessage -> msg.sender == Sender.AGENT && msg.durationMs == null } + // Through removeMessages, so a stored version that follows a dropped bubble is relinked. + removeMessages { unfinished(it) && it.text.isBlank() } + val finalized = _messages.value.map { msg -> + if (unfinished(msg)) msg.copy(status = MessageStatus.COMPLETED, durationMs = 0L) else msg } _messages.value = finalized diff --git a/plugins/AI-Core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/viewmodel/PromptEdit.kt b/plugins/AI-Core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/viewmodel/PromptEdit.kt new file mode 100644 index 00000000..d9b56967 --- /dev/null +++ b/plugins/AI-Core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/viewmodel/PromptEdit.kt @@ -0,0 +1,34 @@ +package com.itsaky.androidide.plugins.aicore.viewmodel + +import com.itsaky.androidide.plugins.aicore.models.ChatMessage +import com.itsaky.androidide.plugins.aicore.models.Sender + +/** + * Decides what an edited prompt does: the newest one is rewound and replaced, an older one forks a + * new branch (see [ChatBranches]). + * + * Kept free of Android and of the ViewModel's run state, so the rules the Edit action follows are + * unit-testable on their own. + */ +internal object PromptEdit { + + /** + * Whether [messageId] is the newest user prompt, the one an edit replaces instead of forking. + * + * @param messages the transcript, oldest first. + */ + fun isLatestPrompt(messages: List, messageId: String): Boolean = + messages.lastOrNull { it.sender == Sender.USER }?.id == messageId + + /** + * The transcript as it stood just before the newest prompt [messageId] was sent. + * + * @param messages the transcript, oldest first. + * @param messageId the prompt being edited. + * @return every message before it, or null when [messageId] is not the newest prompt. + */ + fun messagesBefore(messages: List, messageId: String): List? { + if (!isLatestPrompt(messages, messageId)) return null + return messages.subList(0, messages.indexOfFirst { it.id == messageId }).toList() + } +} diff --git a/plugins/AI-Core/src/main/res/drawable/bg_user_message_editing.xml b/plugins/AI-Core/src/main/res/drawable/bg_user_message_editing.xml new file mode 100644 index 00000000..a0109fe8 --- /dev/null +++ b/plugins/AI-Core/src/main/res/drawable/bg_user_message_editing.xml @@ -0,0 +1,10 @@ + + + + + + + diff --git a/plugins/AI-Core/src/main/res/drawable/ic_chevron_left.xml b/plugins/AI-Core/src/main/res/drawable/ic_chevron_left.xml new file mode 100644 index 00000000..f026f569 --- /dev/null +++ b/plugins/AI-Core/src/main/res/drawable/ic_chevron_left.xml @@ -0,0 +1,10 @@ + + + diff --git a/plugins/AI-Core/src/main/res/drawable/ic_chevron_right.xml b/plugins/AI-Core/src/main/res/drawable/ic_chevron_right.xml new file mode 100644 index 00000000..33afd42d --- /dev/null +++ b/plugins/AI-Core/src/main/res/drawable/ic_chevron_right.xml @@ -0,0 +1,10 @@ + + + diff --git a/plugins/AI-Core/src/main/res/layout/fragment_chat.xml b/plugins/AI-Core/src/main/res/layout/fragment_chat.xml index 06009a9c..237cf884 100644 --- a/plugins/AI-Core/src/main/res/layout/fragment_chat.xml +++ b/plugins/AI-Core/src/main/res/layout/fragment_chat.xml @@ -153,6 +153,46 @@ android:tint="@color/plugin_on_surface_variant" android:visibility="gone" /> + + + + + + + + + + + + + + + + + + + + + + + diff --git a/plugins/AI-Core/src/main/res/values-night/colors.xml b/plugins/AI-Core/src/main/res/values-night/colors.xml index 064f46ae..5d2f598e 100644 --- a/plugins/AI-Core/src/main/res/values-night/colors.xml +++ b/plugins/AI-Core/src/main/res/values-night/colors.xml @@ -30,6 +30,9 @@ #86EFAC #14532D + + #60A5FA + #A3A3A3 #525252 diff --git a/plugins/AI-Core/src/main/res/values/colors.xml b/plugins/AI-Core/src/main/res/values/colors.xml index 0a4f103e..545df3ca 100644 --- a/plugins/AI-Core/src/main/res/values/colors.xml +++ b/plugins/AI-Core/src/main/res/values/colors.xml @@ -35,6 +35,9 @@ #15803D #DCFCE7 + + #2563EB + #737373 #A3A3A3 diff --git a/plugins/AI-Core/src/main/res/values/dimens.xml b/plugins/AI-Core/src/main/res/values/dimens.xml index bbee9d40..3f9c384f 100644 --- a/plugins/AI-Core/src/main/res/values/dimens.xml +++ b/plugins/AI-Core/src/main/res/values/dimens.xml @@ -51,9 +51,13 @@ 2dp 40dp 11dp + + 16dp 14dp 10dp 18dp + + 1.5dp 56dp 20dp diff --git a/plugins/AI-Core/src/main/res/values/integers.xml b/plugins/AI-Core/src/main/res/values/integers.xml index bb535acf..81f0d4ec 100644 --- a/plugins/AI-Core/src/main/res/values/integers.xml +++ b/plugins/AI-Core/src/main/res/values/integers.xml @@ -22,4 +22,6 @@ 4 + + 900 diff --git a/plugins/AI-Core/src/main/res/values/strings.xml b/plugins/AI-Core/src/main/res/values/strings.xml index 2353087d..9693639b 100644 --- a/plugins/AI-Core/src/main/res/values/strings.xml +++ b/plugins/AI-Core/src/main/res/values/strings.xml @@ -16,6 +16,16 @@ Show the whole message Show less of the message Copy message + Edit and resend this message + Editing message + Editing message · the original is kept as a version + Cancel editing + Previous version + Next version + + %1$d/%2$d + + Version %1$d of %2$d Chats, settings and a new chat @@ -94,6 +104,9 @@ Chat cleared Please enter a message Copied + Stop the agent before editing this message + This message is no longer in the chat + Stop the agent before switching versions Generating… Checking the answer… diff --git a/plugins/AI-Core/src/test/kotlin/com/itsaky/androidide/plugins/aicore/viewmodel/ChatBranchesTest.kt b/plugins/AI-Core/src/test/kotlin/com/itsaky/androidide/plugins/aicore/viewmodel/ChatBranchesTest.kt new file mode 100644 index 00000000..19230dcb --- /dev/null +++ b/plugins/AI-Core/src/test/kotlin/com/itsaky/androidide/plugins/aicore/viewmodel/ChatBranchesTest.kt @@ -0,0 +1,137 @@ +package com.itsaky.androidide.plugins.aicore.viewmodel + +import com.itsaky.androidide.plugins.aicore.models.ChatMessage +import com.itsaky.androidide.plugins.aicore.models.ChatSession +import com.itsaky.androidide.plugins.aicore.models.Sender +import org.junit.Assert.assertEquals +import org.junit.Assert.assertNull +import org.junit.Assert.assertSame +import org.junit.Assert.assertTrue +import org.junit.Test + +/** How editing an older prompt forks a chat into versions, and how switching between them works. */ +class ChatBranchesTest { + + private val session = ChatSession( + id = "s1", + messages = listOf( + message("u1", Sender.USER, 1), + message("a1", Sender.AGENT, 2), + message("u2", Sender.USER, 3), + message("t2", Sender.TOOL, 4), + message("a2", Sender.AGENT, 5), + ), + ) + + @Test + fun givenAnUnforkedChat_whenAskingForVersions_thenThereAreNone() { + assertTrue(ChatBranches.positions(session).isEmpty()) + } + + @Test + fun givenAnOlderPrompt_whenForking_thenOnlyTheMessagesAboveItStayOnScreen() { + val forked = ChatBranches.fork(session, "u2")!! + + assertEquals(listOf("u1", "a1"), forked.messages.map { it.id }) + assertEquals(listOf("u2", "t2", "a2"), forked.otherBranches!!.map { it.id }) + assertEquals(listOf("a1", "u2", "t2"), forked.otherBranches!!.map { it.parentId }) + } + + @Test + fun givenAnAgentMessage_whenForking_thenNothingForks() { + assertNull(ChatBranches.fork(session, "a1")) + } + + @Test + fun givenAForkAnsweredByANewPrompt_whenAskingForVersions_thenItIsTheSecondOfTwo() { + val edited = withNewBranch(ChatBranches.fork(session, "u2")!!, "u2b", 10) + + assertEquals(mapOf("u2b" to ChatBranches.Position(1, 2)), ChatBranches.positions(edited)) + assertEquals("u2", ChatBranches.versionId(edited, "u2b", -1)) + assertNull(ChatBranches.versionId(edited, "u2b", 1)) + } + + @Test + fun givenAFork_whenSwitchingBack_thenTheOriginalBranchReturnsWhole() { + val edited = withNewBranch(ChatBranches.fork(session, "u2")!!, "u2b", 10) + + val switched = ChatBranches.switchTo(edited, "u2b", "u2")!! + + assertEquals(listOf("u1", "a1", "u2", "t2", "a2"), switched.messages.map { it.id }) + assertTrue(switched.messages.all { it.parentId == null }) + assertEquals(setOf("u2b", "a2b"), switched.otherBranches!!.map { it.id }.toSet()) + assertEquals(mapOf("u2" to ChatBranches.Position(0, 2)), ChatBranches.positions(switched)) + } + + @Test + fun givenTheFirstPromptForked_whenAskingForVersions_thenTheChatStartHasTwo() { + val edited = withNewBranch(ChatBranches.fork(session, "u1")!!, "u1b", 10) + + assertEquals(listOf("u1b", "a1b"), edited.messages.map { it.id }) + assertEquals(mapOf("u1b" to ChatBranches.Position(1, 2)), ChatBranches.positions(edited)) + } + + @Test + fun givenANestedFork_whenSwitchingAwayAndBack_thenTheVersionLastOnScreenReturns() { + // Fork at u2, answer it with u2b, then fork again at u1 and answer with u1b. + val atU2 = withNewBranch(ChatBranches.fork(session, "u2")!!, "u2b", 10) + val atU1 = withNewBranch(ChatBranches.fork(atU2, "u1")!!, "u1b", 20) + + val back = ChatBranches.switchTo(atU1, "u1b", "u1")!! + + // u2b was on screen under u1 when it was left, not the older u2. + assertEquals(listOf("u1", "a1", "u2b", "a2b"), back.messages.map { it.id }) + assertEquals( + mapOf("u1" to ChatBranches.Position(0, 2), "u2b" to ChatBranches.Position(1, 2)), + ChatBranches.positions(back), + ) + } + + @Test + fun givenAMessageThatIsNotAVersion_whenSwitchingToIt_thenNothingChanges() { + val edited = withNewBranch(ChatBranches.fork(session, "u2")!!, "u2b", 10) + + assertNull(ChatBranches.switchTo(edited, "u2b", "t2")) + } + + @Test + fun givenAVersionFollowingARemovedMessage_whenRemovingIt_thenTheVersionFollowsTheMessageAbove() { + val withNotice = session.copy( + messages = listOf(session.messages[0], session.messages[1], message("n1", Sender.SYSTEM, 3)) + + session.messages.drop(2), + ) + val edited = withNewBranch(ChatBranches.fork(withNotice, "u2")!!, "u2b", 10) + + val trimmed = ChatBranches.removeFromScreen(edited) { it.id == "n1" } + + assertEquals(listOf("u1", "a1", "u2b", "a2b"), trimmed.messages.map { it.id }) + assertEquals("a1", trimmed.otherBranches!!.single { it.id == "u2" }.parentId) + assertEquals("u2", trimmed.selectedBranches!!["a1"]) + assertEquals(mapOf("u2b" to ChatBranches.Position(1, 2)), ChatBranches.positions(trimmed)) + } + + @Test + fun givenNothingToRemove_whenRemovingFromScreen_thenTheSessionIsUnchanged() { + assertSame(session, ChatBranches.removeFromScreen(session) { false }) + } + + @Test + fun givenTwoStoredBranchesSentInTheSameMillisecond_whenSwitching_thenTheOrderIsStable() { + // u2b and u2c share a timestamp, so the id alone must decide which counts as newer. + val atB = withNewBranch(ChatBranches.fork(session, "u2")!!, "u2b", 10) + val atC = withNewBranch(ChatBranches.fork(atB, "u2b")!!, "u2c", 10) + + assertEquals(ChatBranches.Position(2, 3), ChatBranches.positions(atC)["u2c"]) + assertEquals("u2b", ChatBranches.versionId(atC, "u2c", -1)) + } + + /** [forked] with a new prompt [promptId] and its reply appended on screen, as a run adds them. */ + private fun withNewBranch(forked: ChatSession, promptId: String, timestamp: Long) = forked.copy( + messages = forked.messages + + message(promptId, Sender.USER, timestamp) + + message(promptId.replaceFirst('u', 'a'), Sender.AGENT, timestamp + 1), + ) + + private fun message(id: String, sender: Sender, timestamp: Long) = + ChatMessage(id = id, text = "text of $id", sender = sender, timestamp = timestamp) +} diff --git a/plugins/AI-Core/src/test/kotlin/com/itsaky/androidide/plugins/aicore/viewmodel/ChatViewModelPromptEditTest.kt b/plugins/AI-Core/src/test/kotlin/com/itsaky/androidide/plugins/aicore/viewmodel/ChatViewModelPromptEditTest.kt new file mode 100644 index 00000000..c3882c1d --- /dev/null +++ b/plugins/AI-Core/src/test/kotlin/com/itsaky/androidide/plugins/aicore/viewmodel/ChatViewModelPromptEditTest.kt @@ -0,0 +1,440 @@ +package com.itsaky.androidide.plugins.aicore.viewmodel + +import android.content.Context +import android.content.SharedPreferences +import com.itsaky.androidide.plugins.aicore.backends.AiBackend +import com.itsaky.androidide.plugins.aicore.models.AgentState +import com.itsaky.androidide.plugins.aicore.models.Sender +import com.itsaky.androidide.plugins.aicore.prompt.config.DirectoryPromptConfigSource +import com.itsaky.androidide.plugins.aicore.prompt.config.DirectoryPromptConfigSource.Companion.SHIPPED_ROOT +import com.itsaky.androidide.plugins.aicore.prompt.config.sharedPromptConfig +import com.itsaky.androidide.plugins.services.LlmInferenceService +import com.itsaky.androidide.plugins.services.LlmInferenceService.ChatMessage.Role +import com.itsaky.androidide.plugins.services.SharedServices +import io.mockk.every +import io.mockk.mockk +import java.io.File +import java.util.concurrent.CompletableFuture +import java.util.concurrent.ConcurrentHashMap +import kotlinx.coroutines.CoroutineScope +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.ExperimentalCoroutinesApi +import kotlinx.coroutines.SupervisorJob +import kotlinx.coroutines.cancel +import kotlinx.coroutines.launch +import kotlinx.coroutines.runBlocking +import kotlinx.coroutines.test.UnconfinedTestDispatcher +import kotlinx.coroutines.test.resetMain +import kotlinx.coroutines.test.setMain +import org.junit.After +import org.junit.Assert.assertEquals +import org.junit.Assert.assertFalse +import org.junit.Assert.assertTrue +import org.junit.Before +import org.junit.Test + +/** Stand-in project namespace; the digest itself is ProjectKey's business, not this file's. */ +private const val TEST_PROJECT_KEY = "a1b2c3d4e5f60718" +private const val SESSIONS_KEY = "chat_sessions_$TEST_PROJECT_KEY" + +/** Upper bound on waiting for work the ViewModel does on its own dispatchers. */ +private const val WAIT_MS = 5_000L + +/** A reply the agent loop reads as a finished turn: the terminal respond call. */ +private const val RESPOND_REPLY = + """{"tool":"respond","args":{"message":"done"}}""" + +/** + * Editing the newest prompt rewinds the conversation to just before it and runs the edited text + * (ADFA-6276): the transcript, the model's history and the stored session all lose the original, + * and nothing is discarded while a run is in flight or when the send is refused. Editing an older + * prompt forks instead: the original stays as a version, and each version's history is its own. + */ +@OptIn(ExperimentalCoroutinesApi::class) +class ChatViewModelPromptEditTest { + + private lateinit var context: Context + private lateinit var sharedPreferences: SharedPreferences + private lateinit var editor: SharedPreferences.Editor + + /** Backs the preferences mock; concurrent because the persist scope writes from its own thread. */ + private val stored = ConcurrentHashMap() + + private val configScope = CoroutineScope(SupervisorJob() + Dispatchers.Default) + + @Before + fun setUp() { + // ChatViewModel's stateIn() calls run on viewModelScope, i.e. Dispatchers.Main. + Dispatchers.setMain(UnconfinedTestDispatcher()) + SharedServices.clear() + context = mockk(relaxed = true) + sharedPreferences = mockk(relaxed = true) + editor = mockk(relaxed = true) + every { context.getSharedPreferences(any(), any()) } returns sharedPreferences + every { sharedPreferences.getString(any(), any()) } answers { + stored[firstArg()] ?: secondArg() + } + every { sharedPreferences.edit() } returns editor + every { editor.putString(any(), any()) } answers { + val key = firstArg() + val value = secondArg() + if (value == null) stored.remove(key) else stored[key] = value + editor + } + // A run renders its prompts from this config, which the plugin loads on activation. + runBlocking { sharedPromptConfig.preload(configScope, DirectoryPromptConfigSource(SHIPPED_ROOT)).await() } + } + + @After + fun tearDown() { + SharedServices.clear() + Dispatchers.resetMain() + sharedPromptConfig.clear() + configScope.cancel() + } + + @Test + fun givenARestoredChat_whenRewindingToTheNewestPrompt_thenTranscriptAndHistoryEndBeforeIt() { + seedTwoExchanges() + val viewModel = newViewModel() + + assertTrue(viewModel.rewindTo("u2")) + + assertEquals(listOf("u1", "a1"), viewModel.messages.value.map { it.id }) + assertEquals(listOf("u1", "a1"), viewModel.sessions.value.single { it.id == "s1" }.messages.map { it.id }) + assertEquals(listOf(Role.USER, Role.ASSISTANT), viewModel.history.value.map { it.role }) + assertEquals(listOf("first prompt", "first reply"), viewModel.history.value.map { it.content }) + } + + @Test + fun givenAnOlderPrompt_whenRewindingToIt_thenNothingChanges() { + seedTwoExchanges() + val viewModel = newViewModel() + + assertFalse(viewModel.rewindTo("u1")) + + assertEquals(listOf("u1", "a1", "u2", "a2"), viewModel.messages.value.map { it.id }) + assertEquals(4, viewModel.history.value.size) + } + + @Test + fun givenARewind_whenTheChatIsReopened_thenTheDiscardedPromptIsGone() { + seedTwoExchanges() + newViewModel().rewindTo("u2") + + // The write runs on the persist scope's own thread. + awaitCondition { "second prompt" !in stored.getValue(SESSIONS_KEY) } + val reopened = newViewModel() + + assertEquals(listOf("u1", "a1"), reopened.messages.value.map { it.id }) + } + + @Test + fun givenNoBackend_whenEditingTheNewestPrompt_thenItIsRefusedAndNothingIsDiscarded() { + seedTwoExchanges() + val viewModel = newViewModel() + + val result = viewModel.editPrompt("u2", "second prompt, fixed") + + assertEquals(ChatViewModel.EditResult.REFUSED, result) + assertTrue(viewModel.messages.value.any { it.id == "u2" }) + assertEquals(4, viewModel.history.value.size) + } + + @Test + fun givenARunInProgress_whenEditingItsPrompt_thenTheEditIsBlocked() { + // A backend that never answers, so the run stays in flight. + val viewModel = newViewModelWithBackend { } + assertTrue(viewModel.sendMessage("build the app")) + val prompt = awaitUserMessage(viewModel, "build the app") + + val result = viewModel.editPrompt(prompt, "build and run the app") + + assertEquals(ChatViewModel.EditResult.BUSY, result) + assertFalse(viewModel.rewindTo(prompt)) + assertTrue(viewModel.messages.value.any { it.id == prompt }) + viewModel.stopProcessing(reason = "test") + } + + @Test + fun givenAFinishedRun_whenEditingItsPrompt_thenTheNextRunStartsFromBeforeIt() { + val viewModel = newViewModelWithBackend { callback -> + callback.onComplete(LlmInferenceService.LlmResponse.success(RESPOND_REPLY, 4, 10)) + } + assertTrue(viewModel.sendMessage("build teh app")) + val prompt = awaitUserMessage(viewModel, "build teh app") + awaitCondition { viewModel.history.value.size == 2 && viewModel.agentState.value == AgentState.Idle } + + val result = viewModel.editPrompt(prompt, "build the app") + + assertEquals(ChatViewModel.EditResult.STARTED, result) + awaitCondition { + viewModel.history.value.firstOrNull()?.content == "build the app" && + viewModel.agentState.value == AgentState.Idle + } + assertEquals(listOf(Role.USER, Role.ASSISTANT), viewModel.history.value.map { it.role }) + assertEquals( + listOf("build the app"), + viewModel.messages.value.filter { it.sender == Sender.USER }.map { it.text }, + ) + } + + @Test + fun givenTwoFinishedExchanges_whenEditingTheFirstPrompt_thenTheChatForksAndTheModelSeesOnlyTheNewBranch() { + val viewModel = newViewModelWithBackend { callback -> + callback.onComplete(LlmInferenceService.LlmResponse.success(RESPOND_REPLY, 4, 10)) + } + val first = sendAndAwait(viewModel, "build the app") + sendAndAwait(viewModel, "now run the tests") + + val result = viewModel.editPrompt(first, "build the release app") + + assertEquals(ChatViewModel.EditResult.STARTED, result) + awaitCondition { + viewModel.history.value.size == 2 && viewModel.agentState.value == AgentState.Idle + } + assertEquals( + listOf("build the release app"), + viewModel.messages.value.filter { it.sender == Sender.USER }.map { it.text }, + ) + assertEquals("build the release app", viewModel.history.value.first().content) + val edited = viewModel.messages.value.first { it.sender == Sender.USER }.id + assertEquals(ChatBranches.Position(1, 2), viewModel.promptVersions.value[edited]) + } + + @Test + fun givenAForkedChat_whenSwitchingToTheOriginal_thenItsWholeBranchAndHistoryComeBack() { + val viewModel = newViewModelWithBackend { callback -> + callback.onComplete(LlmInferenceService.LlmResponse.success(RESPOND_REPLY, 4, 10)) + } + val first = sendAndAwait(viewModel, "build the app") + sendAndAwait(viewModel, "now run the tests") + viewModel.editPrompt(first, "build the release app") + awaitCondition { + viewModel.history.value.size == 2 && viewModel.agentState.value == AgentState.Idle + } + val edited = viewModel.messages.value.first { it.sender == Sender.USER }.id + + assertTrue(viewModel.switchPromptVersion(edited, -1)) + + assertEquals( + listOf("build the app", "now run the tests"), + viewModel.messages.value.filter { it.sender == Sender.USER }.map { it.text }, + ) + assertTrue(viewModel.history.value.none { it.content == "build the release app" }) + assertEquals(ChatBranches.Position(0, 2), viewModel.promptVersions.value[first]) + assertFalse(viewModel.switchPromptVersion(first, -1)) + } + + @Test + fun givenAForkedChat_whenTheChatIsReopened_thenBothVersionsAreStillThere() { + val viewModel = newViewModelWithBackend { callback -> + callback.onComplete(LlmInferenceService.LlmResponse.success(RESPOND_REPLY, 4, 10)) + } + val first = sendAndAwait(viewModel, "build the app") + sendAndAwait(viewModel, "now run the tests") + viewModel.editPrompt(first, "build the release app") + awaitCondition { + viewModel.history.value.size == 2 && viewModel.agentState.value == AgentState.Idle + } + viewModel.persistState() + awaitCondition { "now run the tests" in stored.getValue(SESSIONS_KEY) } + + val reopened = newViewModel() + + val edited = reopened.messages.value.first { it.sender == Sender.USER } + assertEquals("build the release app", edited.text) + assertEquals(ChatBranches.Position(1, 2), reopened.promptVersions.value[edited.id]) + } + + @Test + fun givenAForkedChat_whenClearingIt_thenNoVersionIsLeft() { + val viewModel = newViewModelWithBackend { callback -> + callback.onComplete(LlmInferenceService.LlmResponse.success(RESPOND_REPLY, 4, 10)) + } + val first = sendAndAwait(viewModel, "build the app") + sendAndAwait(viewModel, "now run the tests") + viewModel.editPrompt(first, "build the release app") + awaitCondition { viewModel.agentState.value == AgentState.Idle && viewModel.history.value.size == 2 } + + viewModel.clearMessages() + + assertTrue(viewModel.sessions.value.all { it.otherBranches.isNullOrEmpty() }) + assertTrue(viewModel.promptVersions.value.isEmpty()) + } + + @Test + fun givenAnAgentMessage_whenEditingIt_thenItIsNotEditable() { + seedTwoExchanges() + val viewModel = newViewModel() + + assertEquals(ChatViewModel.EditResult.NOT_EDITABLE, viewModel.editPrompt("a1", "anything")) + assertEquals(4, viewModel.messages.value.size) + } + + @Test + fun givenAComposerThatRestoresItsDraftOnTheCut_whenEditing_thenThePromptKeepsTheEditsFiles() { + val viewModel = newViewModelWithBackend { callback -> + callback.onComplete(LlmInferenceService.LlmResponse.success(RESPOND_REPLY, 4, 10)) + } + val prompt = sendAndAwait(viewModel, "build the app") + viewModel.setContextFiles(listOf(File("/project/Edited.kt"))) + // As ChatFragment does: once the edited prompt is gone, the draft's files come back. + val composer = CoroutineScope(Dispatchers.Main).launch { + viewModel.messages.collect { messages -> + if (messages.none { it.id == prompt }) viewModel.setContextFiles(listOf(File("/project/Draft.kt"))) + } + } + + assertEquals(ChatViewModel.EditResult.STARTED, viewModel.editPrompt(prompt, "build the release app")) + + composer.cancel() + val edited = viewModel.messages.value.single { it.sender == Sender.USER } + assertEquals(listOf(File("/project/Edited.kt").absolutePath), edited.contextFiles) + awaitCondition { viewModel.agentState.value == AgentState.Idle } + } + + @Test + fun givenTheNewestPromptIsAVersion_whenEditingIt_thenItIsReplacedAndTheOriginalVersionStays() { + val viewModel = newViewModelWithBackend { callback -> + callback.onComplete(LlmInferenceService.LlmResponse.success(RESPOND_REPLY, 4, 10)) + } + val first = sendAndAwait(viewModel, "build the app") + sendAndAwait(viewModel, "now run the tests") + viewModel.editPrompt(first, "build the release app") + awaitCondition { viewModel.agentState.value == AgentState.Idle && viewModel.history.value.size == 2 } + val version = viewModel.messages.value.single { it.sender == Sender.USER }.id + + assertEquals(ChatViewModel.EditResult.STARTED, viewModel.editPrompt(version, "build the debug app")) + awaitCondition { viewModel.agentState.value == AgentState.Idle && viewModel.history.value.size == 2 } + + val replacement = viewModel.messages.value.single { it.sender == Sender.USER } + assertEquals("build the debug app", replacement.text) + assertEquals(ChatBranches.Position(1, 2), viewModel.promptVersions.value[replacement.id]) + val stored = viewModel.sessions.value.single().otherBranches.orEmpty() + assertTrue(stored.none { it.id == version }) + assertTrue(viewModel.switchPromptVersion(replacement.id, -1)) + assertEquals(first, viewModel.messages.value.first().id) + } + + @Test + fun givenAForkedChat_whenSwitchingPastEitherEnd_thenNothingChanges() { + val viewModel = newViewModelWithBackend { callback -> + callback.onComplete(LlmInferenceService.LlmResponse.success(RESPOND_REPLY, 4, 10)) + } + val first = sendAndAwait(viewModel, "build the app") + sendAndAwait(viewModel, "now run the tests") + viewModel.editPrompt(first, "build the release app") + awaitCondition { viewModel.agentState.value == AgentState.Idle && viewModel.history.value.size == 2 } + val edited = viewModel.messages.value.single { it.sender == Sender.USER }.id + + assertFalse(viewModel.switchPromptVersion(edited, 1)) + assertFalse(viewModel.switchPromptVersion(edited, -2)) + assertEquals(edited, viewModel.messages.value.first().id) + } + + @Test + fun givenARunInProgress_whenSwitchingVersions_thenTheSwitchIsBlocked() { + var answer = true + val viewModel = newViewModelWithBackend { callback -> + if (answer) callback.onComplete(LlmInferenceService.LlmResponse.success(RESPOND_REPLY, 4, 10)) + } + val first = sendAndAwait(viewModel, "build the app") + sendAndAwait(viewModel, "now run the tests") + answer = false + + viewModel.editPrompt(first, "build the release app") + val edited = viewModel.messages.value.single { it.sender == Sender.USER }.id + + assertFalse(viewModel.switchPromptVersion(edited, -1)) + assertEquals(edited, viewModel.messages.value.first().id) + viewModel.stopProcessing(reason = "test") + } + + @Test + fun givenAnEditStoppedAtOnce_whenTheRunEnds_thenThePromptHoldsTheForkAndTheChatStillSends() { + var answer = true + val viewModel = newViewModelWithBackend { callback -> + if (answer) callback.onComplete(LlmInferenceService.LlmResponse.success(RESPOND_REPLY, 4, 10)) + } + val first = sendAndAwait(viewModel, "build the app") + sendAndAwait(viewModel, "now run the tests") + answer = false + + viewModel.editPrompt(first, "build the release app") + viewModel.stopProcessing(reason = "test") + + val edited = viewModel.messages.value.single { it.sender == Sender.USER } + assertEquals("build the release app", edited.text) + assertEquals(ChatBranches.Position(1, 2), viewModel.promptVersions.value[edited.id]) + answer = true + awaitCondition { viewModel.sendMessage("and run it") } + } + + /** Sends [text], waits for its run to finish, and returns the prompt's message id. */ + private fun sendAndAwait(viewModel: ChatViewModel, text: String): String { + val before = viewModel.history.value.size + assertTrue(viewModel.sendMessage(text)) + val id = awaitUserMessage(viewModel, text) + awaitCondition { + viewModel.history.value.size == before + 2 && viewModel.agentState.value == AgentState.Idle + } + return id + } + + private fun newViewModel() = ChatViewModel { null }.apply { + initializeStorage(context, TEST_PROJECT_KEY) + } + + /** + * A ViewModel whose selected backend is ready and generates by calling [reply] with the + * turn's callback. + */ + private fun newViewModelWithBackend( + reply: (LlmInferenceService.ToolStreamCallback) -> Unit, + ): ChatViewModel { + val backend = mockk(relaxed = true) + every { backend.id } returns AiBackend.DEFAULT_ID + every { backend.name } returns "Local LLM" + every { backend.isAvailable } returns true + val service = mockk(relaxed = true) + every { service.availableBackends } returns listOf(backend) + every { service.getBackend(any()) } returns backend + // The title request after a first reply; refused, so the chat simply keeps its prompt. + every { service.generateCompletion(any(), any()) } returns + CompletableFuture.completedFuture(LlmInferenceService.LlmResponse.failure("no titles")) + every { service.generateStreamingWithTools(any(), any(), any(), any(), any()) } answers { + reply(arg(4)) + } + SharedServices.register(LlmInferenceService::class.java, service) + return newViewModel().apply { + checkBackendAvailability() + awaitCondition { isBackendAvailable.value } + } + } + + /** @return the id of the user message carrying [text], once the run has posted it. */ + private fun awaitUserMessage(viewModel: ChatViewModel, text: String): String { + awaitCondition { viewModel.messages.value.any { it.sender == Sender.USER && it.text == text } } + return viewModel.messages.value.first { it.sender == Sender.USER && it.text == text }.id + } + + private fun awaitCondition(condition: () -> Boolean) { + val deadline = System.currentTimeMillis() + WAIT_MS + while (!condition()) { + check(System.currentTimeMillis() < deadline) { "condition not met within ${WAIT_MS}ms" } + Thread.sleep(10) + } + } + + private fun seedTwoExchanges() { + stored[SESSIONS_KEY] = + """[{"id":"s1","createdAt":1000,"projectKey":"$TEST_PROJECT_KEY","messages":[""" + + """{"id":"u1","text":"first prompt","sender":"USER","status":"SENT","timestamp":1},""" + + """{"id":"a1","text":"first reply","sender":"AGENT","status":"COMPLETED",""" + + """"timestamp":2,"durationMs":1200},""" + + """{"id":"u2","text":"second prompt","sender":"USER","status":"SENT","timestamp":3},""" + + """{"id":"a2","text":"second reply","sender":"AGENT","status":"COMPLETED",""" + + """"timestamp":4,"durationMs":1200}]}]""" + } +} diff --git a/plugins/AI-Core/src/test/kotlin/com/itsaky/androidide/plugins/aicore/viewmodel/PromptEditTest.kt b/plugins/AI-Core/src/test/kotlin/com/itsaky/androidide/plugins/aicore/viewmodel/PromptEditTest.kt new file mode 100644 index 00000000..470db556 --- /dev/null +++ b/plugins/AI-Core/src/test/kotlin/com/itsaky/androidide/plugins/aicore/viewmodel/PromptEditTest.kt @@ -0,0 +1,66 @@ +package com.itsaky.androidide.plugins.aicore.viewmodel + +import com.itsaky.androidide.plugins.aicore.models.ChatMessage +import com.itsaky.androidide.plugins.aicore.models.Sender +import org.junit.Assert.assertEquals +import org.junit.Assert.assertFalse +import org.junit.Assert.assertNull +import org.junit.Assert.assertTrue +import org.junit.Test + +/** Which prompt an edit replaces rather than forks, and what a rewind to it keeps (ADFA-6276). */ +class PromptEditTest { + + private val transcript = listOf( + message("u1", Sender.USER), + message("a1", Sender.AGENT), + message("u2", Sender.USER), + message("t2", Sender.TOOL), + message("a2", Sender.AGENT), + message("s2", Sender.SYSTEM), + ) + + @Test + fun givenTheNewestUserPrompt_whenAskingIfItIsTheLatest_thenItIs() { + assertTrue(PromptEdit.isLatestPrompt(transcript, "u2")) + } + + @Test + fun givenAnOlderPrompt_whenAskingIfItIsTheLatest_thenItIsNot() { + assertFalse(PromptEdit.isLatestPrompt(transcript, "u1")) + } + + @Test + fun givenNoUserPrompt_whenAskingIfAMessageIsTheLatest_thenNoneIs() { + val agentOnly = listOf(message("s1", Sender.SYSTEM), message("a1", Sender.AGENT)) + + assertFalse(PromptEdit.isLatestPrompt(agentOnly, "a1")) + } + + @Test + fun givenTheNewestPrompt_whenCuttingBeforeIt_thenItAndEverythingAfterItAreDropped() { + val kept = PromptEdit.messagesBefore(transcript, "u2") + + assertEquals(listOf("u1", "a1"), kept?.map { it.id }) + } + + @Test + fun givenAnOlderPrompt_whenCuttingBeforeIt_thenNothingIsCut() { + assertNull(PromptEdit.messagesBefore(transcript, "u1")) + } + + @Test + fun givenAnAgentMessage_whenCuttingBeforeIt_thenNothingIsCut() { + assertNull(PromptEdit.messagesBefore(transcript, "a2")) + } + + @Test + fun givenTheOnlyPrompt_whenCuttingBeforeIt_thenTheTranscriptIsEmpty() { + val single = listOf(message("u1", Sender.USER), message("a1", Sender.AGENT)) + + assertEquals(emptyList(), PromptEdit.messagesBefore(single, "u1")) + } + + private fun message(id: String, sender: Sender) = + ChatMessage(id = id, text = "text of $id", sender = sender) +}