From 37d9b672a412a1c5a9dd731df09267516e8183fa Mon Sep 17 00:00:00 2001 From: Hiep Le <69354317+hieptl@users.noreply.github.com> Date: Fri, 30 Jan 2026 15:54:51 +0700 Subject: [PATCH] fix(frontend): prevent duplicate API calls when sub-conversation becomes ready (planning agent) (#12673) --- .../chat/change-agent-button.test.tsx | 159 +++++++++++ ...use-sub-conversation-task-polling.test.tsx | 253 +++++------------- .../features/chat/change-agent-button.tsx | 36 ++- .../use-sub-conversation-task-polling.ts | 50 +--- 4 files changed, 273 insertions(+), 225 deletions(-) create mode 100644 frontend/__tests__/components/features/chat/change-agent-button.test.tsx diff --git a/frontend/__tests__/components/features/chat/change-agent-button.test.tsx b/frontend/__tests__/components/features/chat/change-agent-button.test.tsx new file mode 100644 index 0000000000..fdc1ba1197 --- /dev/null +++ b/frontend/__tests__/components/features/chat/change-agent-button.test.tsx @@ -0,0 +1,159 @@ +import { screen, waitFor } from "@testing-library/react"; +import { describe, expect, it, vi, beforeEach, afterEach } from "vitest"; +import { QueryClient } from "@tanstack/react-query"; +import { ChangeAgentButton } from "#/components/features/chat/change-agent-button"; +import { renderWithProviders } from "../../../../test-utils"; +import { useConversationStore } from "#/stores/conversation-store"; + +// Mock feature flag to enable planning agent +vi.mock("#/utils/feature-flags", () => ({ + USE_PLANNING_AGENT: () => true, +})); + +// Mock WebSocket status +vi.mock("#/hooks/use-unified-websocket-status", () => ({ + useUnifiedWebSocketStatus: () => "CONNECTED", +})); + +// Mock agent state +vi.mock("#/hooks/use-agent-state", () => ({ + useAgentState: () => ({ curAgentState: "IDLE" }), +})); + +// Track invalidateQueries calls +const mockInvalidateQueries = vi.fn(); + +// Mock react-query to track invalidateQueries calls +vi.mock("@tanstack/react-query", async () => { + const actual = await vi.importActual("@tanstack/react-query"); + return { + ...actual, + useQueryClient: () => ({ + invalidateQueries: mockInvalidateQueries, + }), + }; +}); + +// Mock the active conversation hook +const mockConversationData = { + conversation_id: "parent-conversation-123", + sub_conversation_ids: [], +}; + +vi.mock("#/hooks/query/use-active-conversation", () => ({ + useActiveConversation: () => ({ + data: mockConversationData, + isFetched: true, + refetch: vi.fn(), + }), +})); + +// Mock the sub-conversation task polling hook to control task status +const mockTaskPollingResult = { + task: null as any, + taskStatus: undefined as string | undefined, + taskDetail: null, + taskError: null, + isLoadingTask: false, + subConversationId: undefined as string | undefined, +}; + +vi.mock("#/hooks/query/use-sub-conversation-task-polling", () => ({ + useSubConversationTaskPolling: () => mockTaskPollingResult, +})); + +// Mock the handle plan click hook +vi.mock("#/hooks/use-handle-plan-click", () => ({ + useHandlePlanClick: () => ({ + handlePlanClick: vi.fn(), + isCreatingConversation: false, + }), +})); + +describe("ChangeAgentButton - Cache Invalidation", () => { + beforeEach(() => { + vi.clearAllMocks(); + // Reset store state + useConversationStore.setState({ + conversationMode: "code", + subConversationTaskId: null, + }); + // Reset mock task polling result + mockTaskPollingResult.taskStatus = undefined; + mockTaskPollingResult.subConversationId = undefined; + }); + + afterEach(() => { + vi.clearAllMocks(); + }); + + it("should invalidate parent conversation cache exactly once when task becomes READY", async () => { + // Arrange - Set up a task ID in the store + useConversationStore.setState({ + subConversationTaskId: "task-456", + }); + + // Simulate task becoming READY + mockTaskPollingResult.taskStatus = "READY"; + mockTaskPollingResult.subConversationId = "sub-conversation-789"; + + // Act - Render the component + renderWithProviders(); + + // Assert - Cache should be invalidated exactly once + await waitFor(() => { + expect(mockInvalidateQueries).toHaveBeenCalledTimes(1); + }); + + expect(mockInvalidateQueries).toHaveBeenCalledWith({ + queryKey: ["user", "conversation", "parent-conversation-123"], + }); + }); + + it("should not invalidate cache when task status is not READY", async () => { + // Arrange - Set up a task ID with WORKING status + useConversationStore.setState({ + subConversationTaskId: "task-456", + }); + + mockTaskPollingResult.taskStatus = "WORKING"; + mockTaskPollingResult.subConversationId = undefined; + + // Act + renderWithProviders(); + + // Assert - Wait a bit then verify no invalidation occurred + await new Promise((resolve) => { + setTimeout(resolve, 100); + }); + expect(mockInvalidateQueries).not.toHaveBeenCalled(); + }); + + it("should not invalidate cache when there is no subConversationTaskId", async () => { + // Arrange - No task ID set + useConversationStore.setState({ + subConversationTaskId: null, + }); + + mockTaskPollingResult.taskStatus = "READY"; + mockTaskPollingResult.subConversationId = "sub-conversation-789"; + + // Act + renderWithProviders(); + + // Assert + await new Promise((resolve) => { + setTimeout(resolve, 100); + }); + expect(mockInvalidateQueries).not.toHaveBeenCalled(); + }); + + it("should render the button when planning agent feature is enabled", () => { + // Arrange & Act + renderWithProviders(); + + // Assert + const button = screen.getByRole("button"); + expect(button).toBeInTheDocument(); + }); +}); diff --git a/frontend/__tests__/hooks/query/use-sub-conversation-task-polling.test.tsx b/frontend/__tests__/hooks/query/use-sub-conversation-task-polling.test.tsx index 5dc0dfe76d..2a382bdf86 100644 --- a/frontend/__tests__/hooks/query/use-sub-conversation-task-polling.test.tsx +++ b/frontend/__tests__/hooks/query/use-sub-conversation-task-polling.test.tsx @@ -1,221 +1,112 @@ -import { describe, it, expect, afterEach, vi, beforeEach } from "vitest"; -import React from "react"; import { renderHook, waitFor } from "@testing-library/react"; +import { describe, expect, it, vi, beforeEach, afterEach } from "vitest"; import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; -import { useSubConversationTaskPolling } from "#/hooks/query/use-sub-conversation-task-polling"; +import React from "react"; import V1ConversationService from "#/api/conversation-service/v1-conversation-service.api"; -import { setConversationState } from "#/utils/conversation-local-storage"; -import { useConversationStore } from "#/stores/conversation-store"; +import { useSubConversationTaskPolling } from "#/hooks/query/use-sub-conversation-task-polling"; import type { V1AppConversationStartTask } from "#/api/conversation-service/v1-conversation-service.types"; -// Mock dependencies -vi.mock("#/api/conversation-service/v1-conversation-service.api"); -vi.mock("#/utils/conversation-local-storage"); -vi.mock("#/stores/conversation-store"); +// Mock the underlying service +vi.mock("#/api/conversation-service/v1-conversation-service.api", () => ({ + default: { + getStartTask: vi.fn(), + }, +})); -const mockSetSubConversationTaskId = vi.fn(); -const mockInvalidateQueries = vi.fn(); +describe("useSubConversationTaskPolling", () => { + let queryClient: QueryClient; -// Helper function to create properly typed mock return values -function asMockReturnValue(value: Partial): T { - return value as T; -} + const createWrapper = () => { + queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + }, + }, + }); -function makeTask( - status: V1AppConversationStartTask["status"], - appConversationId: string | null = null, -): V1AppConversationStartTask { - return { + return ({ children }: { children: React.ReactNode }) => ( + {children} + ); + }; + + const createMockTask = ( + status: V1AppConversationStartTask["status"], + appConversationId: string | null = null, + ): V1AppConversationStartTask => ({ id: "task-123", - created_by_user_id: "user-123", + created_by_user_id: "user-1", status, detail: null, app_conversation_id: appConversationId, sandbox_id: null, agent_server_url: null, - request: { - agent_type: "plan", - parent_conversation_id: "parent-conv-123", - }, + request: {}, created_at: new Date().toISOString(), updated_at: new Date().toISOString(), - }; -} - -function createQueryClient() { - return new QueryClient({ - defaultOptions: { - queries: { - retry: false, - }, - }, }); -} -function wrapper({ children }: { children: React.ReactNode }) { - const queryClient = createQueryClient(); - queryClient.invalidateQueries = mockInvalidateQueries; - return ( - {children} - ); -} - -describe("useSubConversationTaskPolling", () => { beforeEach(() => { vi.clearAllMocks(); - vi.mocked(useConversationStore).mockReturnValue( - asMockReturnValue>({ - setSubConversationTaskId: mockSetSubConversationTaskId, - }), - ); }); afterEach(() => { - vi.clearAllMocks(); + queryClient?.clear(); }); - describe("when task status is READY", () => { - it("clears subConversationTaskId from localStorage and store when task completes successfully", async () => { - const parentConversationId = "parent-conv-123"; - const taskId = "task-123"; - const appConversationId = "sub-conv-456"; + it("should return task status when task is READY", async () => { + // Arrange + const mockTask = createMockTask("READY", "sub-conversation-123"); + vi.mocked(V1ConversationService.getStartTask).mockResolvedValue(mockTask); - const readyTask = makeTask("READY", appConversationId); - vi.mocked(V1ConversationService.getStartTask).mockResolvedValue( - readyTask, - ); + // Act + const { result } = renderHook( + () => + useSubConversationTaskPolling("task-123", "parent-conversation-456"), + { wrapper: createWrapper() }, + ); - renderHook( - () => useSubConversationTaskPolling(taskId, parentConversationId), - { wrapper }, - ); - - await waitFor(() => { - expect(setConversationState).toHaveBeenCalledWith( - parentConversationId, - { subConversationTaskId: null }, - ); - }); - - expect(mockSetSubConversationTaskId).toHaveBeenCalledWith(null); - expect(mockInvalidateQueries).toHaveBeenCalledWith({ - queryKey: ["user", "conversation", parentConversationId], - }); - }); - - it("invalidates parent conversation cache when task is READY", async () => { - const parentConversationId = "parent-conv-123"; - const taskId = "task-123"; - const appConversationId = "sub-conv-456"; - - const readyTask = makeTask("READY", appConversationId); - vi.mocked(V1ConversationService.getStartTask).mockResolvedValue( - readyTask, - ); - - renderHook( - () => useSubConversationTaskPolling(taskId, parentConversationId), - { wrapper }, - ); - - await waitFor(() => { - expect(mockInvalidateQueries).toHaveBeenCalled(); - }); - - expect(mockInvalidateQueries).toHaveBeenCalledWith({ - queryKey: ["user", "conversation", parentConversationId], - }); + // Assert + await waitFor(() => { + expect(result.current.taskStatus).toBe("READY"); }); + expect(result.current.subConversationId).toBe("sub-conversation-123"); + expect(V1ConversationService.getStartTask).toHaveBeenCalledWith("task-123"); }); - describe("when task status is ERROR", () => { - it("clears subConversationTaskId from localStorage and store on error", async () => { - const parentConversationId = "parent-conv-123"; - const taskId = "task-123"; + it("should not poll when taskId is null", async () => { + // Arrange + vi.mocked(V1ConversationService.getStartTask).mockResolvedValue(null); - const errorTask = makeTask("ERROR", null); - vi.mocked(V1ConversationService.getStartTask).mockResolvedValue( - errorTask, - ); + // Act + const { result } = renderHook( + () => useSubConversationTaskPolling(null, "parent-conversation-456"), + { wrapper: createWrapper() }, + ); - renderHook( - () => useSubConversationTaskPolling(taskId, parentConversationId), - { wrapper }, - ); - - await waitFor(() => { - expect(setConversationState).toHaveBeenCalledWith( - parentConversationId, - { subConversationTaskId: null }, - ); - }); - - expect(mockSetSubConversationTaskId).toHaveBeenCalledWith(null); - expect(mockInvalidateQueries).not.toHaveBeenCalled(); + // Assert - wait a bit to ensure no calls are made + await new Promise((resolve) => { + setTimeout(resolve, 100); }); + expect(V1ConversationService.getStartTask).not.toHaveBeenCalled(); + expect(result.current.taskStatus).toBeUndefined(); }); - describe("when task is still in progress", () => { - it("does not clear subConversationTaskId when task status is WORKING", async () => { - const parentConversationId = "parent-conv-123"; - const taskId = "task-123"; + it("should not poll when parentConversationId is null", async () => { + // Arrange + vi.mocked(V1ConversationService.getStartTask).mockResolvedValue(null); - const workingTask = makeTask("WORKING", null); - vi.mocked(V1ConversationService.getStartTask).mockResolvedValue( - workingTask, - ); + // Act + const { result } = renderHook( + () => useSubConversationTaskPolling("task-123", null), + { wrapper: createWrapper() }, + ); - renderHook( - () => useSubConversationTaskPolling(taskId, parentConversationId), - { wrapper }, - ); - - // Wait a bit to ensure useEffect has run - await waitFor(() => { - expect(V1ConversationService.getStartTask).toHaveBeenCalled(); - }); - - // Should not clear anything while task is in progress - expect(setConversationState).not.toHaveBeenCalled(); - expect(mockSetSubConversationTaskId).not.toHaveBeenCalled(); - }); - - it("does not clear subConversationTaskId when task status is WAITING_FOR_SANDBOX", async () => { - const parentConversationId = "parent-conv-123"; - const taskId = "task-123"; - - const waitingTask = makeTask("WAITING_FOR_SANDBOX", null); - vi.mocked(V1ConversationService.getStartTask).mockResolvedValue( - waitingTask, - ); - - renderHook( - () => useSubConversationTaskPolling(taskId, parentConversationId), - { wrapper }, - ); - - await waitFor(() => { - expect(V1ConversationService.getStartTask).toHaveBeenCalled(); - }); - - expect(setConversationState).not.toHaveBeenCalled(); - expect(mockSetSubConversationTaskId).not.toHaveBeenCalled(); - }); - }); - - describe("when parentConversationId is null", () => { - it("does not clear subConversationTaskId or invalidate queries", () => { - const taskId = "task-123"; - - renderHook(() => useSubConversationTaskPolling(taskId, null), { - wrapper, - }); - - // Query is disabled when parentConversationId is null, so getStartTask won't be called - expect(V1ConversationService.getStartTask).not.toHaveBeenCalled(); - expect(setConversationState).not.toHaveBeenCalled(); - expect(mockSetSubConversationTaskId).not.toHaveBeenCalled(); - expect(mockInvalidateQueries).not.toHaveBeenCalled(); + // Assert - wait a bit to ensure no calls are made + await new Promise((resolve) => { + setTimeout(resolve, 100); }); + expect(V1ConversationService.getStartTask).not.toHaveBeenCalled(); + expect(result.current.taskStatus).toBeUndefined(); }); }); diff --git a/frontend/src/components/features/chat/change-agent-button.tsx b/frontend/src/components/features/chat/change-agent-button.tsx index 28de891db9..f95e2b6a4f 100644 --- a/frontend/src/components/features/chat/change-agent-button.tsx +++ b/frontend/src/components/features/chat/change-agent-button.tsx @@ -1,5 +1,6 @@ -import React, { useMemo, useEffect, useState } from "react"; +import React, { useMemo, useEffect, useState, useRef } from "react"; import { useTranslation } from "react-i18next"; +import { useQueryClient } from "@tanstack/react-query"; import { Typography } from "#/ui/typography"; import { I18nKey } from "#/i18n/declaration"; import CodeTagIcon from "#/icons/code-tag.svg?react"; @@ -36,12 +37,41 @@ export function ChangeAgentButton() { const { data: conversation } = useActiveConversation(); - // Poll sub-conversation task and invalidate parent conversation when ready - useSubConversationTaskPolling( + const queryClient = useQueryClient(); + + // Track the last invalidated task ID to prevent duplicate invalidations + const lastInvalidatedTaskIdRef = useRef(null); + + // Poll sub-conversation task status + const { taskStatus, subConversationId } = useSubConversationTaskPolling( subConversationTaskId, conversation?.conversation_id || null, ); + // Invalidate parent conversation cache when task is ready (only once per task) + useEffect(() => { + if ( + taskStatus === "READY" && + subConversationId && + conversation?.conversation_id && + subConversationTaskId && + lastInvalidatedTaskIdRef.current !== subConversationTaskId + ) { + // Mark this task as invalidated to prevent duplicate calls + lastInvalidatedTaskIdRef.current = subConversationTaskId; + // Invalidate the parent conversation to refetch with updated sub_conversation_ids + queryClient.invalidateQueries({ + queryKey: ["user", "conversation", conversation.conversation_id], + }); + } + }, [ + taskStatus, + subConversationId, + conversation?.conversation_id, + subConversationTaskId, + queryClient, + ]); + // Get handlePlanClick and isCreatingConversation from custom hook const { handlePlanClick, isCreatingConversation } = useHandlePlanClick(); diff --git a/frontend/src/hooks/query/use-sub-conversation-task-polling.ts b/frontend/src/hooks/query/use-sub-conversation-task-polling.ts index 96ab17ea31..c97fb729fa 100644 --- a/frontend/src/hooks/query/use-sub-conversation-task-polling.ts +++ b/frontend/src/hooks/query/use-sub-conversation-task-polling.ts @@ -1,31 +1,28 @@ -import { useEffect } from "react"; -import { useQuery, useQueryClient } from "@tanstack/react-query"; +import { useQuery } from "@tanstack/react-query"; import V1ConversationService from "#/api/conversation-service/v1-conversation-service.api"; -import { setConversationState } from "#/utils/conversation-local-storage"; -import { useConversationStore } from "#/stores/conversation-store"; /** - * Hook that polls V1 sub-conversation start tasks and invalidates parent conversation cache when ready. + * Hook that polls V1 sub-conversation start tasks. * * This hook: * - Polls the V1 start task API every 3 seconds until status is READY or ERROR - * - Automatically invalidates the parent conversation cache when the task becomes READY * - Exposes task status and details for UI components to show loading states and errors * + * Note: This hook does NOT invalidate the parent conversation cache. The component + * that initiates the sub-conversation creation should handle cache invalidation + * to ensure it only happens once. + * * Use case: - * - When creating a sub-conversation (e.g., plan mode), track the task and refresh parent conversation - * data once the sub-conversation is ready + * - When creating a sub-conversation (e.g., plan mode), track the task status + * for UI loading states * * @param taskId - The task ID to poll (from createConversation response) - * @param parentConversationId - The parent conversation ID to invalidate when ready + * @param parentConversationId - The parent conversation ID (used to enable polling) */ export const useSubConversationTaskPolling = ( taskId: string | null, parentConversationId: string | null, ) => { - const queryClient = useQueryClient(); - const { setSubConversationTaskId } = useConversationStore(); - // Poll the task if we have both taskId and parentConversationId const taskQuery = useQuery({ queryKey: ["sub-conversation-task", taskId], @@ -49,35 +46,6 @@ export const useSubConversationTaskPolling = ( retry: false, }); - // Invalidate parent conversation cache when task is ready and clear localStorage - useEffect(() => { - const task = taskQuery.data; - if (!parentConversationId) return; - - if (task?.status === "READY" && task.app_conversation_id) { - // Invalidate the parent conversation to refetch with updated sub_conversation_ids - queryClient.invalidateQueries({ - queryKey: ["user", "conversation", parentConversationId], - }); - // Clear the task ID from localStorage and store since task completed successfully - setConversationState(parentConversationId, { - subConversationTaskId: null, - }); - setSubConversationTaskId(null); - } else if (task?.status === "ERROR") { - // Clear the task ID from localStorage and store on error so user can retry - setConversationState(parentConversationId, { - subConversationTaskId: null, - }); - setSubConversationTaskId(null); - } - }, [ - taskQuery.data, - parentConversationId, - queryClient, - setSubConversationTaskId, - ]); - return { task: taskQuery.data, taskStatus: taskQuery.data?.status,