Allow editing the last stopped Agent prompt

browser_tests/tests/agent/agentPanel.spec.ts

@@ -192,6 +192,45 @@ test.describe('In-App Agent panel', { tag: '@cloud' }, () => {
     expect(track.scrollbarColor).toMatch(/rgba\(0, 0, 0, 0\)$/)
   })
 
+  test('edits and resubmits the last prompt after stopping its turn', async ({
+    comfyPage,
+    postedMessages,
+    getWebSocket
+  }) => {
+    const page = comfyPage.page
+    await page.getByRole('button', { name: OPEN_AGENT_LABEL }).click()
+
+    const panel = page.locator('#agent-panel-root')
+    const composer = panel.getByRole('textbox', { name: /^Describe ideas/ })
+    const originalPrompt = 'Build a rainy city at night'
+    const revisedPrompt = 'Build a rainy city at sunrise'
+
+    await composer.fill(originalPrompt)
+    await panel.getByRole('button', { name: enMessages.agent.send }).click()
+    await expect.poll(() => postedMessages.length).toBe(1)
+
+    await expect(
+      panel.getByRole('button', { name: enMessages.g.edit })
+    ).toHaveCount(0)
+    await panel.getByRole('button', { name: enMessages.agent.stop }).click()
+    await expect(
+      panel.getByRole('button', { name: enMessages.g.edit })
+    ).toHaveCount(0)
+
+    pushEvent(await getWebSocket(), MESSAGE_DONE_EVENT)
+    const editButton = panel.getByRole('button', { name: enMessages.g.edit })
+    await expect(editButton).toHaveCount(1)
+    await editButton.click()
+
+    await expect(composer).toHaveValue(originalPrompt)
+    await expect(composer).toBeFocused()
+
+    await composer.fill(revisedPrompt)
+    await panel.getByRole('button', { name: enMessages.agent.send }).click()
+    await expect.poll(() => postedMessages.length).toBe(2)
+    expect(postedMessages[1]).toContain(revisedPrompt)
+  })
+
   test('applies a draft_patch graph to the canvas', async ({
     comfyPage,
     postedMessages,

browser_tests/tests/agent/agentPanelMocks.ts

@@ -4,6 +4,7 @@ import { comfyPageFixture } from '@e2e/fixtures/ComfyPage'
 
 import type { RemoteConfig } from '@/platform/remoteConfig/types'
 import type {
+  AgentCancelAccepted,
   AgentDraftSnapshot,
   AgentTurnAccepted,
   AgentWsEvent,
@@ -25,6 +26,8 @@ const TURN_ACCEPTED: AgentTurnAccepted = {
   workflow_id: WORKFLOW_ID
 }
 
+const CANCEL_ACCEPTED: AgentCancelAccepted = { status: 'cancelling' }
+
 const DRAFT_GRAPH: ComfyWorkflowJSON = {
   version: 0.4,
   last_node_id: 2,
@@ -261,15 +264,26 @@ async function mockAgentBoot(
     const request = route.request()
     if (request.method() === 'POST') {
       postedMessages.push(request.postData() ?? '')
+      const accepted: AgentTurnAccepted = {
+        ...TURN_ACCEPTED,
+        message_id:
+          postedMessages.length === 1
+            ? TURN_ID
+            : `${TURN_ID}-${postedMessages.length}`
+      }
       return route.fulfill({
         status: 202,
         contentType: 'application/json',
-        body: JSON.stringify(TURN_ACCEPTED)
+        body: JSON.stringify(accepted)
       })
     }
     return route.fulfill(jsonRoute([]))
   })
 
+  await page.route('**/api/agent/threads/*/messages/*/cancel', (route: Route) =>
+    route.fulfill(jsonRoute(CANCEL_ACCEPTED))
+  )
+
   await page.route('**/api/agent/draft**', (r) =>
     r.fulfill(jsonRoute(DRAFT_SNAPSHOT))
   )

src/workbench/extensions/agent/AgentPanelRoot.vue

@@ -340,6 +340,7 @@ const {
   start,
   stop,
   entries,
+  editableTurnId,
   isStreaming,
   status,
   notices,
@@ -1161,6 +1162,7 @@ function onPanelDrop(event: DragEvent): void {
     <AgentPanel
       ref="panelRef"
       :entries
+      :editable-turn-id="editableTurnId"
       :user-name="userName"
       :streaming="isStreaming"
       :submitting="isSending || status === 'thinking'"

src/workbench/extensions/agent/components/agent/AgentPanel.test.ts

@@ -4,6 +4,7 @@ import { createPinia, setActivePinia } from 'pinia'
 import { beforeEach, describe, expect, it } from 'vitest'
 
 import { i18n } from '@/i18n'
+import type { TurnId } from '../../schemas/agentApiSchema'
 
 import AgentPanel from './AgentPanel.vue'
 
@@ -88,4 +89,36 @@ describe('AgentPanel', () => {
 
     expect(textarea).not.toHaveFocus()
   })
+
+  it('replaces and focuses the composer draft when editing the eligible prompt', async () => {
+    const user = userEvent.setup()
+    const pinia = createPinia()
+    setActivePinia(pinia)
+    const prompt = 'Generate a yellow duck with a hockey mask'
+    const { emitted } = render(AgentPanel, {
+      props: {
+        editableTurnId: 'msg-1' as TurnId,
+        entries: [{ id: 'msg-1' as TurnId, role: 'user', text: prompt }],
+        historyGroups
+      },
+      global: {
+        plugins: [pinia, i18n],
+        directives: { tooltip: {} },
+        stubs: { WorkflowSelectorChip: true }
+      }
+    })
+    const textarea = screen.getByRole('textbox')
+    await user.type(textarea, 'unfinished draft')
+
+    await user.click(screen.getByRole('button', { name: 'Edit' }))
+
+    expect(textarea).toHaveValue(prompt)
+    expect(textarea).toHaveFocus()
+
+    await user.clear(textarea)
+    await user.type(textarea, 'Generate a yellow duck at sunrise')
+    await user.click(screen.getByRole('button', { name: 'Send' }))
+
+    expect(emitted().send[0]).toEqual(['Generate a yellow duck at sunrise', []])
+  })
 })

src/workbench/extensions/agent/components/agent/AgentPanel.vue

@@ -14,6 +14,7 @@ import { buildAgentTooltipConfig } from '@/composables/useTooltipConfig'
 import type { AssetItem } from '@/platform/assets/schemas/assetSchema'
 
 import type { ActiveTab } from '../../types/activeTab'
+import type { TurnId } from '../../schemas/agentApiSchema'
 import type { ComposerAttachment } from '../../composables/agent/useComposer'
 import type { SelectedNode } from '../../composables/agent/useCanvasSelection'
 import type { ConversationEntry } from '../../stores/agent/agentConversationStore'
@@ -43,7 +44,8 @@ const {
   getMentionAssets = async () => [],
   sessionId = null,
   customTitle,
-  historyGroups
+  historyGroups,
+  editableTurnId = null
 } = defineProps<{
   entries: ConversationEntry[]
   userName?: string
@@ -61,6 +63,7 @@ const {
   sessionId?: string | null
   customTitle?: string
   historyGroups: HistoryGroups
+  editableTurnId?: TurnId | null
 }>()
 const emit = defineEmits<{
   send: [text: string, attachments: ComposerAttachment[]]
@@ -275,6 +278,8 @@ defineExpose({ addAttachment, updateAttachment, removeAttachment })
         <ConversationView
           v-else
           :entries="entries"
+          :editable-turn-id="editableTurnId"
+          @edit-prompt="composerRef?.replaceDraft($event)"
           @feedback="(id, vote) => emit('feedback', id, vote)"
         />
       </div>

src/workbench/extensions/agent/components/agent/Composer.vue

@@ -296,8 +296,14 @@ function insert(text: string): void {
   textareaRef.value?.focus()
 }
 
+function replaceDraft(text: string): void {
+  composer.draft.value = text
+  textareaRef.value?.focus()
+}
+
 defineExpose({
   insert,
+  replaceDraft,
   addAttachment: composer.addAttachment,
   updateAttachment: composer.updateAttachment,
   removeAttachment: composer.removeAttachment

src/workbench/extensions/agent/components/agent/ConversationView.vue

@@ -8,15 +8,18 @@ import { buildAgentTooltipConfig } from '@/composables/useTooltipConfig'
 import { cn } from '@comfyorg/tailwind-utils'
 
 import type { ConversationEntry } from '../../stores/agent/agentConversationStore'
+import type { TurnId } from '../../schemas/agentApiSchema'
 
 import AgentMessage from './message/AgentMessage.vue'
 import UserMessage from './message/UserMessage.vue'
 
-const { entries } = defineProps<{
+const { entries, editableTurnId = null } = defineProps<{
   entries: ConversationEntry[]
+  editableTurnId?: TurnId | null
 }>()
 const emit = defineEmits<{
   feedback: [turnId: string, vote: 'up' | 'down' | null]
+  editPrompt: [text: string]
 }>()
 
 const { t } = useI18n()
@@ -77,6 +80,8 @@ watch(
               :text="entry.text"
               :attachments="entry.attachments"
               :tags="entry.tags"
+              :editable="entry.id === editableTurnId"
+              @edit="emit('editPrompt', $event)"
             />
             <AgentMessage
               v-else

src/workbench/extensions/agent/components/agent/message/UserMessage.test.ts

@@ -28,6 +28,7 @@ function renderMessage(props: {
   text: string
   attachments?: { name: string; previewUrl?: string }[]
   tags?: string[]
+  editable?: boolean
 }) {
   return render(UserMessage, { props, global: { plugins: [i18n] } })
 }
@@ -80,6 +81,29 @@ describe('UserMessage', () => {
     expect(clipboard.copy).toHaveBeenCalledWith('make it cinematic')
   })
 
+  it('offers an accessible edit action only when the prompt is editable', async () => {
+    const user = userEvent.setup()
+    const prompt = 'make it cinematic'
+    const { emitted } = renderMessage({ text: prompt, editable: true })
+
+    const editButton = screen.getByRole('button', { name: t('g.edit') })
+    await user.hover(editButton)
+    expect(
+      await screen.findByRole('tooltip', { hidden: true })
+    ).toHaveTextContent(t('g.edit'))
+    await user.click(editButton)
+
+    expect(emitted().edit).toEqual([[prompt]])
+  })
+
+  it('does not offer edit for a settled prompt without edit eligibility', () => {
+    renderMessage({ text: 'make it cinematic' })
+
+    expect(
+      screen.queryByRole('button', { name: t('g.edit') })
+    ).not.toBeInTheDocument()
+  })
+
   it('offers no copy action on an attachment-only message', () => {
     renderMessage({ text: '', attachments: [{ name: 'clip.bin' }] })
 

src/workbench/extensions/agent/components/agent/message/UserMessage.vue

@@ -9,11 +9,16 @@ import AgentTooltip from '../AgentTooltip.vue'
 const {
   text,
   attachments = [],
-  tags = []
+  tags = [],
+  editable = false
 } = defineProps<{
   text: string
   attachments?: UserAttachment[]
   tags?: string[]
+  editable?: boolean
+}>()
+const emit = defineEmits<{
+  edit: [text: string]
 }>()
 
 const { t } = useI18n()
@@ -66,8 +71,18 @@ const { copy, copied } = useClipboard({ copiedDuring: 2000, legacy: true })
     </div>
     <div
       v-if="text"
-      class="text-agent-fg-subtle flex opacity-0 transition-opacity group-hover:opacity-100 focus-within:opacity-100"
+      class="text-agent-fg-subtle pointer-events-none flex opacity-0 transition-opacity group-hover:pointer-events-auto group-hover:opacity-100 focus-within:pointer-events-auto focus-within:opacity-100 touch:pointer-events-auto touch:opacity-100"
     >
+      <AgentTooltip v-if="editable" :label="t('g.edit')">
+        <button
+          type="button"
+          :aria-label="t('g.edit')"
+          class="hover:bg-agent-surface-hover hover:text-agent-fg flex size-6 cursor-pointer items-center justify-center rounded-lg p-1 transition-colors"
+          @click="emit('edit', text)"
+        >
+          <span class="icon-[lucide--pencil] size-3" />
+        </button>
+      </AgentTooltip>
       <AgentTooltip :label="copied ? t('agent.copied') : t('agent.copy')">
         <button
           type="button"

src/workbench/extensions/agent/composables/agent/useAgentSession.test.ts

@@ -375,28 +375,60 @@ describe('useAgentSession (v1 composition root)', () => {
   })
 
   it('(d) stopTurn cancels the active turn; a 409 is swallowed and the socket settles it', async () => {
+    const postMessage = vi
+      .fn<
+        (threadId: string, req: PostMessageInput) => Promise<AgentTurnAccepted>
+      >()
+      .mockResolvedValueOnce({ thread_id: 'th-1', message_id: 'msg-1' })
+      .mockResolvedValueOnce({ thread_id: 'th-1', message_id: 'msg-2' })
     const cancelMessage = vi
       .fn<
         (threadId: string, messageId: string) => Promise<AgentCancelAccepted>
       >()
       .mockRejectedValue(new AgentApiError('already done', 409, undefined))
-    const rest = fakeRest({ cancelMessage })
+    const rest = fakeRest({ cancelMessage, postMessage })
     const { source, emit } = fakeEvents()
     const session = useAgentSession({ rest, events: source })
     session.start()
 
     await session.sendMessage('go')
     emit(delta('msg-1', 'working'))
     expect(session.isStreaming.value).toBe(true)
+    expect(session.editableTurnId.value).toBeNull()
 
     await session.stopTurn()
     expect(cancelMessage).toHaveBeenCalledWith('th-1', 'msg-1')
     expect(session.notices.value).toHaveLength(0)
     expect(session.isStreaming.value).toBe(true)
+    expect(session.editableTurnId.value).toBeNull()
 
     emit(delta('msg-1', ' Stopped at your request.'))
     emit(done('msg-1'))
     expect(session.isStreaming.value).toBe(false)
+    expect(session.editableTurnId.value).toBe('msg-1')
+
+    await session.sendMessage('go revised')
+    expect(session.editableTurnId.value).toBeNull()
+    expect(
+      session.entries.value
+        .filter((entry) => entry.role === 'user')
+        .map((entry) => entry.text)
+    ).toEqual(['go', 'go revised'])
+
+    session.newChat()
+    expect(session.editableTurnId.value).toBeNull()
+  })
+
+  it('(d1) a normally completed turn is not editable', async () => {
+    const { source, emit } = fakeEvents()
+    const session = useAgentSession({ rest: fakeRest(), events: source })
+    session.start()
+
+    await session.sendMessage('go')
+    emit(done('msg-1'))
+
+    expect(session.isStreaming.value).toBe(false)
+    expect(session.editableTurnId.value).toBeNull()
   })
 
   it('(d2) stopTurn rejecting with a network TypeError surfaces a notice, not an unhandled rejection', async () => {
@@ -419,6 +451,7 @@ describe('useAgentSession (v1 composition root)', () => {
     expect(session.notices.value).toEqual([
       { level: 'error', text: 'fetch failed' }
     ])
+    expect(session.editableTurnId.value).toBeNull()
   })
 
   it('(d3) stopTurn before the POST acknowledgement cancels the acknowledged turn once', async () => {