eigent/test/integration/chatStore/deadWorkforce.test.tsx

// ========= Copyright 2025-2026 @ Eigent.ai All Rights Reserved. =========
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ========= Copyright 2025-2026 @ Eigent.ai All Rights Reserved. =========

import { act, renderHook, waitFor } from '@testing-library/react';
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';

// Import proxy mock to enable API mocking
import '../../mocks/proxy.mock';
// Also Mock authStore & sse
import '../../mocks/authStore.mock';
import '../../mocks/sse.mock';

// Import chat store to ensure it's available
import '../../../src/store/chatStore';

import useChatStoreAdapter from '../../../src/hooks/useChatStoreAdapter';
import { useProjectStore } from '../../../src/store/projectStore';
import { mockFetchEventSource } from '../../mocks/sse.mock';

// Helper function for sequential SSE events
const createSSESequence = (events: Array<{ event: any; delay: number }>) => {
  return async (onMessage: (data: any) => void) => {
    for (let i = 0; i < events.length; i++) {
      const { event, delay } = events[i];

      await new Promise<void>((resolve) => {
        setTimeout(() => {
          console.log(`Sending SSE Event ${i + 1}:`, event.step);
          onMessage({
            data: JSON.stringify(event),
          });
          resolve();
        }, delay);
      });
    }
  };
};

// Mock electron IPC
(global as any).ipcRenderer = {
  invoke: vi.fn((channel) => {
    if (channel === 'get-system-language') return Promise.resolve('en');
    if (channel === 'get-browser-port') return Promise.resolve(9222);
    if (channel === 'get-env-path') return Promise.resolve('/path/to/env');
    if (channel === 'mcp-list') return Promise.resolve({});
    if (channel === 'get-file-list') return Promise.resolve([]);
    return Promise.resolve();
  }),
};

// Mock window.electronAPI
Object.defineProperty(window, 'electronAPI', {
  value: {
    uploadLog: vi.fn().mockResolvedValue(undefined),
    // Add other electronAPI methods as needed
  },
  writable: true,
});

describe('Integration Test: Case 2 - same session new chat', () => {
  beforeEach(() => {
    vi.clearAllMocks();

    const { result } = renderHook(() => useProjectStore());
    //Reset projectStore
    result.current.getAllProjects().forEach((project) => {
      result.current.removeProject(project.id);
    });

    //Create initial Project
    const projectId = result.current.createProject(
      'Complete Journey Test',
      'Testing full flow'
    );
    expect(projectId).toBeDefined();

    // 2. Get chatStore (automatically created)
    let chatStore = result.current.getActiveChatStore(projectId)!;
    expect(chatStore).toBeDefined();
    const initiatorTaskId = chatStore.getState().activeTaskId!;
    expect(initiatorTaskId).toBeDefined();
  });

  afterEach(() => {
    vi.clearAllMocks();
  });

  it('Sequential startTask appends chatStores to same project', async () => {
    const { result, rerender } = renderHook(() => useChatStoreAdapter());

    // Setup the events sequence
    const eventSequence = createSSESequence([
      {
        event: {
          step: 'to_sub_tasks',
          data: {
            summary_task: 'Calculator App|Build a simple calculator',
            sub_tasks: [
              { id: 'task-1', content: 'Create UI components', status: '' },
              {
                id: 'task-2',
                content: 'Implement calculator logic',
                status: '',
              },
            ],
          },
        },
        delay: 100,
      },
      {
        event: {
          step: 'end',
          data: "--- Subtask 1760647257350-6021.1 Result ---\nI'm doing well, thank you for asking! How are you today?",
        },
        delay: 400,
      },
    ]);

    // Mock SSE stream with controlled events
    mockFetchEventSource.mockImplementation(
      async (url: string, options: any) => {
        if (options.onmessage) {
          await eventSequence(options.onmessage);
        }
      }
    );

    await act(async () => {
      // Complete Flow Test
      const { chatStore, projectStore: _projectStore } = result.current;
      const initiatorTaskId = chatStore.activeTaskId;

      //User Message to send
      const userMessage = 'Build a calculator app';

      // 6. Start task
      // NOTE: startTask creates a NEW chatStore and switches to it, the old chatStore is no longer active
      await chatStore.startTask(
        initiatorTaskId,
        undefined,
        undefined,
        undefined,
        userMessage
      );

      //Rerender to get the latest chatStore
      rerender();
    });

    // Test 1: Check initial task creation and first SSE event (to_sub_tasks)
    await waitFor(
      () => {
        rerender();
        const { chatStore: newChatStore, projectStore } = result.current;
        expect(newChatStore).toBeDefined();

        let taskId = newChatStore?.activeTaskId!;
        const task = newChatStore?.tasks[taskId];
        expect(taskId).toBeDefined();

        if (task) {
          expect(task.hasMessages).toBe(true);
          expect(task.messages[0].content).toBe('Build a calculator app');
          expect(task.status).toBe('pending');
        }

        const updatedTask = newChatStore?.tasks[taskId];
        expect(updatedTask?.summaryTask).toBe(
          'Calculator App|Build a simple calculator'
        );
        expect(updatedTask?.taskInfo).toHaveLength(3);
        expect(updatedTask?.taskRunning).toHaveLength(3);

        //Two chatStores - first initial
        expect(
          projectStore.getAllChatStores(projectStore.activeProjectId)
        ).toHaveLength(2);
      },
      { timeout: 1000 }
    );

    // Test 2: Check progress event has been processed
    await waitFor(
      () => {
        rerender();
        const { chatStore: newChatStore } = result.current;
        let taskId = newChatStore?.activeTaskId!;
        const task = newChatStore?.tasks[taskId];

        // Check if progress message exists in messages or status updates
        // Adjust this based on how your app handles progress events
        expect(task).toBeDefined();
        console.log('Progress test - task status:', task?.status);
      },
      { timeout: 1500 }
    );

    // Test 3: Rerender until status is "finished"
    await waitFor(
      () => {
        rerender();
        const { chatStore: newChatStore } = result.current;
        let taskId = newChatStore?.activeTaskId!;
        const task = newChatStore?.tasks[taskId];

        if (task) {
          // Check if task is completed or has final result
          // Adjust these assertions based on your app's behavior
          console.log('End test - task status:', task?.status);
          console.log('End test - task messages:', task?.messages?.length);
        }
        expect(task.status).toBe('finished');
        console.log(task);
      },
      { timeout: 2000 }
    );

    //Before starting new chatStore
    const { chatStore, projectStore } = result.current;
    expect(Object.keys(chatStore.tasks)).toHaveLength(1);
    //Initial ChatStore + appendedOne
    expect(
      projectStore.getAllChatStores(projectStore.activeProjectId)
    ).toHaveLength(2);
    //Make all tasks are skipped after end
    chatStore.tasks[chatStore.activeTaskId].taskRunning.forEach((task: any) => {
      expect(task.status).toBe('skipped');
    });
    chatStore.tasks[chatStore.activeTaskId].taskInfo.forEach((task: any) => {
      expect(task.status).toBe('skipped');
    });

    // Test: Start second chat session with different events
    await act(async () => {
      rerender();
      const { chatStore, projectStore: _projectStore } = result.current;
      const initiatorTaskId = chatStore.activeTaskId;

      // Setup different events for second session
      const secondEventSequence = createSSESequence([
        {
          event: {
            step: 'confirmed',
            data: { question: 'how are you?' },
          },
          delay: 100,
        },
        {
          event: {
            step: 'to_sub_tasks',
            data: {
              summary_task: 'Todo App|Build a todo application',
              sub_tasks: [
                { id: 'task-3', content: 'Design todo interface', status: '' },
                { id: 'task-4', content: 'Implement todo logic', status: '' },
              ],
            },
          },
          delay: 200,
        },
        {
          event: {
            step: 'end',
            data: '--- Second Task Result ---\nTodo app planning completed!',
          },
          delay: 300,
        },
        {
          event: {
            step: 'end',
            data: "--- Subtask 1760647257350-6021.1 Result ---\nI'm doing well, thank you for asking! How are you today?",
          },
          delay: 400,
        },
      ]);

      // Update the mock for the second call
      mockFetchEventSource.mockImplementation(
        async (url: string, options: any) => {
          if (options.onmessage) {
            await secondEventSequence(options.onmessage);
          }
        }
      );

      const userMessage = 'Build a todo app';
      await chatStore.startTask(
        initiatorTaskId,
        undefined,
        undefined,
        undefined,
        userMessage
      );
      rerender();
    });

    // Test the second session results
    await waitFor(
      () => {
        rerender();
        const { chatStore: newChatStore, projectStore } = result.current;
        let taskId = newChatStore?.activeTaskId!;
        const task = newChatStore?.tasks[taskId];

        if (task) {
          expect(task.messages[0].content).toBe('Build a todo app');
          // Check if the new summary task is set correctly
          expect(task.summaryTask).toBe('Todo App|Build a todo application');
        }

        // Should now have 3 chat stores (initial + 2 task sessions)
        expect(
          projectStore.getAllChatStores(projectStore.activeProjectId)
        ).toHaveLength(3);
      },
      { timeout: 1500 }
    );

    // CHECK POST End State of Run 2
    await waitFor(
      () => {
        rerender();
        const { chatStore: newChatStore } = result.current;
        let taskId = newChatStore?.activeTaskId!;
        const task = newChatStore?.tasks[taskId];

        if (task) {
          // Check if task is completed or has final result
          // Adjust these assertions based on your app's behavior
          console.log('End test 2 - task status:', task?.status);
          console.log('End test 2 - task messages:', task?.messages?.length);
        }
        expect(task.status).toBe('finished');
        console.log(task);
      },
      { timeout: 2000 }
    );

    //Before starting new chatStore
    const { chatStore: secondChatStore } = result.current;
    expect(Object.keys(secondChatStore.tasks)).toHaveLength(1);
    //Initial ChatStore + appendedOne
    expect(
      projectStore.getAllChatStores(projectStore.activeProjectId)
    ).toHaveLength(3);
    //Make all tasks are skipped after end
    secondChatStore.tasks[secondChatStore.activeTaskId].taskRunning.forEach(
      (task: any) => {
        expect(task.status).toBe('skipped');
      }
    );
    secondChatStore.tasks[secondChatStore.activeTaskId].taskInfo.forEach(
      (task: any) => {
        expect(task.status).toBe('skipped');
      }
    );
  });

  it('should handle individual SSE events with precise timing', async () => {
    const { result, rerender } = renderHook(() => useChatStoreAdapter());

    let messageCallback: ((data: any) => void) | null = null;

    // Mock SSE to capture the callback for manual event triggering
    mockFetchEventSource.mockImplementation((url: string, options: any) => {
      messageCallback = options.onmessage;
      // Don't send any events automatically
    });

    // Start the task
    await act(async () => {
      const { chatStore } = result.current;
      const initiatorTaskId = chatStore.activeTaskId;
      await chatStore.startTask(
        initiatorTaskId,
        undefined,
        undefined,
        undefined,
        'Test manual events'
      );
      rerender();
    });

    // Manually send first event and test result
    await act(async () => {
      messageCallback?.({
        data: JSON.stringify({
          step: 'to_sub_tasks',
          data: {
            summary_task: 'Manual Test|Testing manual event control',
            sub_tasks: [
              { id: 'manual-1', content: 'Manual task 1', status: '' },
            ],
          },
        }),
      });
    });

    // Test first event result
    await waitFor(() => {
      rerender();
      const { chatStore } = result.current;
      let taskId = chatStore?.activeTaskId!;
      const task = chatStore?.tasks[taskId];
      expect(task?.summaryTask).toBe(
        'Manual Test|Testing manual event control'
      );
    });

    // Send second event and test result
    await act(async () => {
      messageCallback?.({
        data: JSON.stringify({
          step: 'progress',
          data: 'Manual progress update',
        }),
      });
    });

    // Test second event result
    await waitFor(() => {
      rerender();
      const { chatStore } = result.current;
      let taskId = chatStore?.activeTaskId!;
      const task = chatStore?.tasks[taskId];
      // Add your specific assertions for progress events here
      expect(task).toBeDefined();
    });

    // Send final event
    await act(async () => {
      messageCallback?.({
        data: JSON.stringify({
          step: 'end',
          data: 'Manual test completed successfully',
        }),
      });
    });

    // Test final state
    await waitFor(() => {
      rerender();
      const { chatStore } = result.current;
      let taskId = chatStore?.activeTaskId!;
      const task = chatStore?.tasks[taskId];
      // Add your specific assertions for end state here
      expect(task).toBeDefined();
    });
  });

  //TODO: Don't let new startTask until newChatStore appended
  it('Parallel startTask calls with separate chatStores (startTask -> wait for append -> startTask)', async () => {
    const { result, rerender } = renderHook(() => useChatStoreAdapter());

    let sseCallCount = 0;
    let firstTaskChatStore: any = null;
    let secondTaskChatStore: any = null;

    // Setup SSE events for the first task
    const firstTaskEventSequence = createSSESequence([
      {
        event: {
          step: 'to_sub_tasks',
          data: {
            summary_task: 'First Task|Build a calculator app',
            sub_tasks: [
              { id: 'first-1', content: 'Create calculator UI', status: '' },
              { id: 'first-2', content: 'Implement calc logic', status: '' },
            ],
          },
        },
        delay: 100,
      },
      {
        event: {
          step: 'end',
          data: '--- First Task Result ---\nCalculator app completed successfully!',
        },
        delay: 300,
      },
    ]);

    // Setup SSE events for the second task
    const secondTaskEventSequence = createSSESequence([
      {
        event: {
          step: 'to_sub_tasks',
          data: {
            summary_task: 'Second Task|Build a todo app',
            sub_tasks: [
              { id: 'second-1', content: 'Create todo UI', status: '' },
              { id: 'second-2', content: 'Implement todo logic', status: '' },
            ],
          },
        },
        delay: 100,
      },
      {
        event: {
          step: 'end',
          data: '--- Second Task Result ---\nTodo app completed successfully!',
        },
        delay: 300,
      },
    ]);

    // Mock SSE to handle sequential calls with different events
    mockFetchEventSource.mockImplementation(
      async (url: string, options: any) => {
        sseCallCount++;
        console.log(`SSE Call #${sseCallCount} initiated`);

        if (sseCallCount === 1) {
          // First task gets first event sequence
          console.log('Processing first task events');
          if (options.onmessage) {
            await firstTaskEventSequence(options.onmessage);
          }
        } else if (sseCallCount === 2) {
          // Second task gets second event sequence
          console.log('Processing second task events');
          if (options.onmessage) {
            await secondTaskEventSequence(options.onmessage);
          }
        }
      }
    );

    // Get initial chatStore reference
    const { chatStore: initialChatStore, projectStore: _projectStore } =
      result.current;
    const initialChatStoreRef = initialChatStore;
    const initiatorTaskId = initialChatStore.activeTaskId;

    // Step 1: Start first task
    await act(async () => {
      const userMessage1 = 'Build a calculator app';
      await initialChatStore.startTask(
        initiatorTaskId,
        undefined,
        undefined,
        undefined,
        userMessage1
      );
      rerender();
    });

    // Verify first task started and chatStore was created
    await waitFor(
      () => {
        rerender();
        const { chatStore: currentChatStore, projectStore } = result.current;

        // Should have 2 chatStores: initial + first task
        const allChatStores = projectStore.getAllChatStores(
          projectStore.activeProjectId
        );
        expect(allChatStores).toHaveLength(2);

        // Current chatStore should be different from initial (new one created)
        expect(currentChatStore).not.toBe(initialChatStoreRef);
        firstTaskChatStore = currentChatStore;

        // Verify first task details
        const activeTaskId = currentChatStore.activeTaskId;
        const activeTask = currentChatStore.tasks[activeTaskId];
        expect(activeTask).toBeDefined();
        expect(activeTask.hasMessages).toBe(true);
        expect(activeTask.messages[0].content).toBe('Build a calculator app');
      },
      { timeout: 1000 }
    );

    // Wait for first task SSE events to process (summary_task)
    await waitFor(
      () => {
        rerender();
        const { chatStore: currentChatStore } = result.current;
        const activeTaskId = currentChatStore.activeTaskId;
        const activeTask = currentChatStore.tasks[activeTaskId];

        // Check that SSE events have been processed
        expect(activeTask.summaryTask).toBe(
          'First Task|Build a calculator app'
        );
        expect(activeTask.taskInfo).toHaveLength(3); // main task + 2 subtasks
        expect(activeTask.taskRunning).toHaveLength(3);

        console.log('First task SSE events processed');
      },
      { timeout: 1500 }
    );

    // Wait for first task to complete
    await waitFor(
      () => {
        rerender();
        const { chatStore: currentChatStore } = result.current;
        const activeTaskId = currentChatStore.activeTaskId;
        const activeTask = currentChatStore.tasks[activeTaskId];

        expect(activeTask.status).toBe('finished');

        console.log('First task completed');
      },
      { timeout: 2000 }
    );

    // Step 2: Wait for project append to complete before starting second task
    await waitFor(
      () => {
        rerender();
        const { projectStore } = result.current;

        // Ensure the project has been properly updated with the appended first chatStore
        const allChatStores = projectStore.getAllChatStores(
          projectStore.activeProjectId
        );
        expect(allChatStores).toHaveLength(2);

        // Verify the active chatStore is properly set and ready for next task
        const activeChatStore = projectStore.getActiveChatStore(
          projectStore.activeProjectId
        );
        expect(activeChatStore).toBeDefined();
        expect(activeChatStore.getState()?.activeTaskId).toBeDefined();

        console.log('Project append completed, ready for second task');
      },
      { timeout: 1000 }
    );

    // Step 3: Start second task on the same chatStore
    await act(async () => {
      rerender();
      const { chatStore: currentChatStore } = result.current;
      const currentInitiatorTaskId = currentChatStore.activeTaskId;

      const userMessage2 = 'Build a todo app';
      await currentChatStore.startTask(
        currentInitiatorTaskId,
        undefined,
        undefined,
        undefined,
        userMessage2
      );
      rerender();
    });

    // Verify second task started and new chatStore was created
    await waitFor(
      () => {
        rerender();
        const { chatStore: currentChatStore, projectStore } = result.current;

        // Should now have 3 chatStores: initial + first task + second task
        const allChatStores = projectStore.getAllChatStores(
          projectStore.activeProjectId
        );
        expect(allChatStores).toHaveLength(3);

        // Current chatStore should be different from first task chatStore
        expect(currentChatStore).not.toBe(firstTaskChatStore);
        secondTaskChatStore = currentChatStore;

        // Verify second task details
        const activeTaskId = currentChatStore.activeTaskId;
        const activeTask = currentChatStore.tasks[activeTaskId];
        expect(activeTask).toBeDefined();
        expect(activeTask.hasMessages).toBe(true);
        expect(activeTask.messages[0].content).toBe('Build a todo app');
      },
      { timeout: 1000 }
    );

    // Wait for second task SSE events to process (summary_task)
    await waitFor(
      () => {
        rerender();
        const { chatStore: currentChatStore } = result.current;
        const activeTaskId = currentChatStore.activeTaskId;
        const activeTask = currentChatStore.tasks[activeTaskId];

        // Check that SSE events have been processed
        expect(activeTask.summaryTask).toBe('Second Task|Build a todo app');
        expect(activeTask.taskInfo).toHaveLength(3); // main task + 2 subtasks
        expect(activeTask.taskRunning).toHaveLength(3);

        console.log('Second task SSE events processed');
      },
      { timeout: 1500 }
    );

    // Wait for second task to complete
    await waitFor(
      () => {
        rerender();
        const { chatStore: currentChatStore } = result.current;
        const activeTaskId = currentChatStore.activeTaskId;
        const activeTask = currentChatStore.tasks[activeTaskId];

        expect(activeTask.status).toBe('finished');

        console.log('Second task completed');
      },
      { timeout: 2000 }
    );

    // Final verification: Both chatStores should have separate states
    const { projectStore: finalProjectStore } = result.current;
    const allFinalChatStores = finalProjectStore.getAllChatStores(
      finalProjectStore.activeProjectId
    );

    // Verify we have 3 separate chatStores with their own states
    expect(allFinalChatStores).toHaveLength(3);
    expect(sseCallCount).toBe(2); // Two separate SSE calls

    // Verify each chatStore has its own task with different content
    expect(firstTaskChatStore).not.toBe(secondTaskChatStore);

    // Get the current state of chatStores from the project store to verify final states
    const [
      _initialChatStoreFromProject,
      firstChatStoreFromProject,
      secondChatStoreFromProject,
    ] = allFinalChatStores;

    // Verify first chatStore state (should be the second in the array after initial)
    const firstTaskId =
      firstChatStoreFromProject.chatStore.getState().activeTaskId;
    const firstTask =
      firstChatStoreFromProject.chatStore.getState().tasks[firstTaskId];
    expect(firstTask.messages[0].content).toBe('Build a calculator app');
    expect(firstTask.summaryTask).toBe('First Task|Build a calculator app');

    // Verify second chatStore state (should be the third in the array)
    const secondTaskId =
      secondChatStoreFromProject.chatStore.getState().activeTaskId;
    const secondTask =
      secondChatStoreFromProject.chatStore.getState().tasks[secondTaskId];
    expect(secondTask.messages[0].content).toBe('Build a todo app');
    expect(secondTask.summaryTask).toBe('Second Task|Build a todo app');

    console.log(
      'Sequential startTask test with separate chatStores completed successfully'
    );
  });
});