Add tests for ChatHistoryWidget and related features (#3921)

browser_tests/tests/chatHistory.spec.ts

@@ -0,0 +1,139 @@
+import { Page, expect } from '@playwright/test'
+
+import { comfyPageFixture as test } from '../fixtures/ComfyPage'
+
+interface ChatHistoryEntry {
+  prompt: string
+  response: string
+  response_id: string
+}
+
+async function renderChatHistory(page: Page, history: ChatHistoryEntry[]) {
+  const nodeId = await page.evaluate(() => window['app'].graph.nodes[0]?.id)
+  // Simulate API sending display_component message
+  await page.evaluate(
+    ({ nodeId, history }) => {
+      const event = new CustomEvent('display_component', {
+        detail: {
+          node_id: nodeId,
+          component: 'ChatHistoryWidget',
+          props: {
+            history: JSON.stringify(history)
+          }
+        }
+      })
+      window['app'].api.dispatchEvent(event)
+      return true
+    },
+    { nodeId, history }
+  )
+
+  return nodeId
+}
+
+test.describe('Chat History Widget', () => {
+  let nodeId: string
+
+  test.beforeEach(async ({ comfyPage }) => {
+    nodeId = await renderChatHistory(comfyPage.page, [
+      { prompt: 'Hello', response: 'World', response_id: '123' }
+    ])
+    // Wait for chat history to be rendered
+    await comfyPage.page.waitForSelector('.pi-pencil')
+  })
+
+  test('displays chat history when receiving display_component message', async ({
+    comfyPage
+  }) => {
+    // Verify the chat history is displayed correctly
+    await expect(comfyPage.page.getByText('Hello')).toBeVisible()
+    await expect(comfyPage.page.getByText('World')).toBeVisible()
+  })
+
+  test('handles message editing interaction', async ({ comfyPage }) => {
+    // Get first node's ID
+    nodeId = await comfyPage.page.evaluate(() => {
+      const node = window['app'].graph.nodes[0]
+
+      // Make sure the node has a prompt widget (for editing functionality)
+      if (!node.widgets) {
+        node.widgets = []
+      }
+
+      // Add a prompt widget if it doesn't exist
+      if (!node.widgets.find((w) => w.name === 'prompt')) {
+        node.widgets.push({
+          name: 'prompt',
+          type: 'text',
+          value: 'Original prompt'
+        })
+      }
+
+      return node.id
+    })
+
+    await renderChatHistory(comfyPage.page, [
+      {
+        prompt: 'Message 1',
+        response: 'Response 1',
+        response_id: '123'
+      },
+      {
+        prompt: 'Message 2',
+        response: 'Response 2',
+        response_id: '456'
+      }
+    ])
+    await comfyPage.page.waitForSelector('.pi-pencil')
+
+    const originalTextAreaInput = await comfyPage.page
+      .getByPlaceholder('text')
+      .nth(1)
+      .inputValue()
+
+    // Click edit button on first message
+    await comfyPage.page.getByLabel('Edit').first().click()
+    await comfyPage.nextFrame()
+
+    // Verify cancel button appears
+    await expect(comfyPage.page.getByLabel('Cancel')).toBeVisible()
+
+    // Click cancel edit
+    await comfyPage.page.getByLabel('Cancel').click()
+
+    // Verify prompt input is restored
+    await expect(comfyPage.page.getByPlaceholder('text').nth(1)).toHaveValue(
+      originalTextAreaInput
+    )
+  })
+
+  test('handles real-time updates to chat history', async ({ comfyPage }) => {
+    // Send initial history
+    await renderChatHistory(comfyPage.page, [
+      {
+        prompt: 'Initial message',
+        response: 'Initial response',
+        response_id: '123'
+      }
+    ])
+    await comfyPage.page.waitForSelector('.pi-pencil')
+
+    // Update history with additional messages
+    await renderChatHistory(comfyPage.page, [
+      {
+        prompt: 'Follow-up',
+        response: 'New response',
+        response_id: '456'
+      }
+    ])
+    await comfyPage.page.waitForSelector('.pi-pencil')
+
+    // Move mouse over the canvas to force update
+    await comfyPage.page.mouse.move(100, 100)
+    await comfyPage.nextFrame()
+
+    // Verify new messages appear
+    await expect(comfyPage.page.getByText('Follow-up')).toBeVisible()
+    await expect(comfyPage.page.getByText('New response')).toBeVisible()
+  })
+})

src/components/graph/widgets/ChatHistoryWidget.vue

@@ -96,8 +96,7 @@ const setPromptInput = (text: string, previousResponseId?: string | null) => {
 }
 
 const handleEdit = (index: number) => {
-  if (!promptInput) return
-
+  promptInput ??= widget?.node.widgets?.find((w) => w.name === 'prompt')
   editIndex.value = index
   const prevResponseId = index === 0 ? 'start' : getPreviousResponseId(index)
   const promptText = parsedHistory.value[index]?.prompt ?? ''

src/composables/node/useNodeChatHistory.ts

@@ -16,9 +16,6 @@ export function useNodeChatHistory(
 ) {
   const chatHistoryWidget = useChatHistoryWidget(options)
 
-  const findChatHistoryWidget = (node: LGraphNode) =>
-    node.widgets?.find((w) => w.name === CHAT_HISTORY_WIDGET_NAME)
-
   const addChatHistoryWidget = (node: LGraphNode) =>
     chatHistoryWidget(node, {
       name: CHAT_HISTORY_WIDGET_NAME,
@@ -30,9 +27,11 @@ export function useNodeChatHistory(
    * @param node The graph node to show the chat history for
    */
   function showChatHistory(node: LGraphNode) {
-    if (!findChatHistoryWidget(node)) {
-      addChatHistoryWidget(node)
-    }
+    // First remove any existing widget
+    removeChatHistory(node)
+
+    // Then add the widget with new history
+    addChatHistoryWidget(node)
     node.setDirtyCanvas?.(true)
   }
 

tests-ui/tests/components/ChatHistoryWidget.spec.ts

@@ -0,0 +1,95 @@
+import { mount } from '@vue/test-utils'
+import { describe, expect, it, vi } from 'vitest'
+import { createI18n } from 'vue-i18n'
+
+import ChatHistoryWidget from '@/components/graph/widgets/ChatHistoryWidget.vue'
+
+const i18n = createI18n({
+  legacy: false,
+  locale: 'en',
+  messages: {
+    en: {
+      g: { edit: 'Edit' },
+      chatHistory: {
+        cancelEdit: 'Cancel edit',
+        cancelEditTooltip: 'Cancel edit'
+      }
+    }
+  }
+})
+
+vi.mock('@/components/graph/widgets/chatHistory/CopyButton.vue', () => ({
+  default: {
+    name: 'CopyButton',
+    template: '<div class="mock-copy-button"></div>',
+    props: ['text']
+  }
+}))
+
+vi.mock('@/components/graph/widgets/chatHistory/ResponseBlurb.vue', () => ({
+  default: {
+    name: 'ResponseBlurb',
+    template: '<div class="mock-response-blurb"><slot /></div>',
+    props: ['text']
+  }
+}))
+
+describe('ChatHistoryWidget.vue', () => {
+  const mockHistory = JSON.stringify([
+    { prompt: 'Test prompt', response: 'Test response', response_id: '123' }
+  ])
+
+  const mountWidget = (props: { history: string; widget?: any }) => {
+    return mount(ChatHistoryWidget, {
+      props,
+      global: {
+        plugins: [i18n],
+        stubs: {
+          Button: {
+            template: '<button><slot /></button>',
+            props: ['icon', 'aria-label']
+          },
+          ScrollPanel: { template: '<div><slot /></div>' }
+        }
+      }
+    })
+  }
+
+  it('renders chat history correctly', () => {
+    const wrapper = mountWidget({ history: mockHistory })
+    expect(wrapper.text()).toContain('Test prompt')
+    expect(wrapper.text()).toContain('Test response')
+  })
+
+  it('handles empty history', () => {
+    const wrapper = mountWidget({ history: '[]' })
+    expect(wrapper.find('.mb-4').exists()).toBe(false)
+  })
+
+  it('edits previous prompts', () => {
+    const mockWidget = {
+      node: { widgets: [{ name: 'prompt', value: '' }] }
+    }
+
+    const wrapper = mountWidget({ history: mockHistory, widget: mockWidget })
+    const vm = wrapper.vm as any
+    vm.handleEdit(0)
+
+    expect(mockWidget.node.widgets[0].value).toContain('Test prompt')
+    expect(mockWidget.node.widgets[0].value).toContain('starting_point_id')
+  })
+
+  it('cancels editing correctly', () => {
+    const mockWidget = {
+      node: { widgets: [{ name: 'prompt', value: 'Original value' }] }
+    }
+
+    const wrapper = mountWidget({ history: mockHistory, widget: mockWidget })
+    const vm = wrapper.vm as any
+
+    vm.handleEdit(0)
+    vm.handleCancelEdit()
+
+    expect(mockWidget.node.widgets[0].value).toBe('Original value')
+  })
+})

tests-ui/tests/composables/useNodeChatHistory.test.ts

@@ -0,0 +1,63 @@
+import { LGraphNode } from '@comfyorg/litegraph'
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+
+import { useNodeChatHistory } from '@/composables/node/useNodeChatHistory'
+
+vi.mock('@/composables/widgets/useChatHistoryWidget', () => ({
+  useChatHistoryWidget: () => {
+    return (node: any, inputSpec: any) => {
+      const widget = {
+        name: inputSpec.name,
+        type: inputSpec.type
+      }
+
+      if (!node.widgets) {
+        node.widgets = []
+      }
+      node.widgets.push(widget)
+
+      return widget
+    }
+  }
+}))
+
+// Mock LGraphNode type
+type MockNode = {
+  widgets: Array<{ name: string; type: string }>
+  setDirtyCanvas: ReturnType<typeof vi.fn>
+  addCustomWidget: ReturnType<typeof vi.fn>
+  [key: string]: any
+}
+
+describe('useNodeChatHistory', () => {
+  const mockNode = {
+    widgets: [],
+    setDirtyCanvas: vi.fn(),
+    addCustomWidget: vi.fn()
+  } as unknown as LGraphNode & MockNode
+
+  beforeEach(() => {
+    mockNode.widgets = []
+    mockNode.setDirtyCanvas.mockClear()
+    mockNode.addCustomWidget.mockClear()
+  })
+
+  it('adds chat history widget to node', () => {
+    const { showChatHistory } = useNodeChatHistory()
+    showChatHistory(mockNode)
+
+    expect(mockNode.widgets.length).toBe(1)
+    expect(mockNode.widgets[0].name).toBe('$$node-chat-history')
+    expect(mockNode.setDirtyCanvas).toHaveBeenCalled()
+  })
+
+  it('removes chat history widget from node', () => {
+    const { showChatHistory, removeChatHistory } = useNodeChatHistory()
+    showChatHistory(mockNode)
+
+    expect(mockNode.widgets.length).toBe(1)
+
+    removeChatHistory(mockNode)
+    expect(mockNode.widgets.length).toBe(0)
+  })
+})

tests-ui/tests/store/executionStore.test.ts

@@ -0,0 +1,82 @@
+import { createPinia, setActivePinia } from 'pinia'
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+
+import { useExecutionStore } from '@/stores/executionStore'
+
+// Remove any previous global types
+declare global {
+  // Empty interface to override any previous declarations
+  interface Window {}
+}
+
+const mockShowChatHistory = vi.fn()
+vi.mock('@/composables/node/useNodeChatHistory', () => ({
+  useNodeChatHistory: () => ({
+    showChatHistory: mockShowChatHistory
+  })
+}))
+
+vi.mock('@/composables/node/useNodeProgressText', () => ({
+  useNodeProgressText: () => ({
+    showTextPreview: vi.fn()
+  })
+}))
+
+// Create a local mock instead of using global to avoid conflicts
+const mockApp = {
+  graph: {
+    getNodeById: vi.fn()
+  }
+}
+
+describe('executionStore - display_component handling', () => {
+  function createDisplayComponentEvent(
+    nodeId: string,
+    component = 'ChatHistoryWidget'
+  ) {
+    return new CustomEvent('display_component', {
+      detail: {
+        node_id: nodeId,
+        component,
+        props: {
+          history: JSON.stringify([{ prompt: 'Test', response: 'Response' }])
+        }
+      }
+    })
+  }
+
+  function handleDisplayComponentMessage(event: CustomEvent) {
+    const { node_id, component } = event.detail
+    const node = mockApp.graph.getNodeById(node_id)
+    if (node && component === 'ChatHistoryWidget') {
+      mockShowChatHistory(node)
+    }
+  }
+
+  beforeEach(() => {
+    setActivePinia(createPinia())
+    useExecutionStore()
+    vi.clearAllMocks()
+  })
+
+  it('handles ChatHistoryWidget display_component messages', () => {
+    const mockNode = { id: '123' }
+    mockApp.graph.getNodeById.mockReturnValue(mockNode)
+
+    const event = createDisplayComponentEvent('123')
+    handleDisplayComponentMessage(event)
+
+    expect(mockApp.graph.getNodeById).toHaveBeenCalledWith('123')
+    expect(mockShowChatHistory).toHaveBeenCalledWith(mockNode)
+  })
+
+  it('does nothing if node is not found', () => {
+    mockApp.graph.getNodeById.mockReturnValue(null)
+
+    const event = createDisplayComponentEvent('non-existent')
+    handleDisplayComponentMessage(event)
+
+    expect(mockApp.graph.getNodeById).toHaveBeenCalledWith('non-existent')
+    expect(mockShowChatHistory).not.toHaveBeenCalled()
+  })
+})