fix(frontend): prevent duplicate API calls when sub-conversation becomes ready (planning agent) (#12673)

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(<ChangeAgentButton />);
+
+    // 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(<ChangeAgentButton />);
+
+    // 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(<ChangeAgentButton />);
+
+    // 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(<ChangeAgentButton />);
+
+    // Assert
+    const button = screen.getByRole("button");
+    expect(button).toBeInTheDocument();
+  });
+});

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<T>(value: Partial<T>): 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 }) => (
+      <QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
+    );
+  };
+
+  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 (
-    <QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
-  );
-}
-
-describe("useSubConversationTaskPolling", () => {
   beforeEach(() => {
     vi.clearAllMocks();
-    vi.mocked(useConversationStore).mockReturnValue(
-      asMockReturnValue<ReturnType<typeof useConversationStore>>({
-        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();
   });
 });

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<string | null>(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();
 

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,